o
    ¹i9  ã                   @   sÔ   d Z ddlZddlZddlZddlmZ ddlmZ ddlmZ ddl	m
Z
 ddlmZmZ ddlmZ dd	lmZ d
dlmZ dd„ ZdedƒedƒdœdœZeƒ Zdd„ Zdd„ ZedeƒZG dd„ deƒZdS )zTest GatewayClienté    N)Údatetime)ÚStringIO)Úpatch)Ú	HTTPError)ÚHTTPRequestÚHTTPResponse)ÚGatewayClient)Úmaybe_futureé   )ÚNotebookTestBasec                 C   s.   g d¢}d|i | ddi dœi}| |i dœ}|S )N)Úpythonz-mÚipykernel_launcherz-fz{connection_file}Úspecr   Úsignal)ÚargvÚenvÚdisplay_nameÚlanguageÚinterrupt_modeÚmetadata)Únamer   Ú	resources© )r   Úargv_stanzaÚspec_stanzaÚkernelspec_stanzar   r   úV/var/www/edux/Edux_v2/venv/lib/python3.10/site-packages/notebook/tests/test_gateway.pyÚgenerate_kernelspec   s   r   Ú	kspec_fooÚ	kspec_bar)r   r   )ÚdefaultÚkernelspecsc                 C   s4   t  ¡  ¡ d }tt ¡ ƒ}|| t|ƒdddœ}|S )zfGenerate a mocked kernel model.  Caller is responsible for adding model to running_kernels dictionary.ÚZÚidler
   )Úidr   Úlast_activityÚexecution_stateÚconnections)r   ÚutcnowÚ	isoformatÚstrÚuuidÚuuid4)r   ÚdtÚ	kernel_idÚmodelr   r   r   Úgenerate_model    s   r0   c                 Ë   sp  d}|d r|d }t dd| i|¤Ž}t| ƒ}| d¡r6|dkr6tt t¡ƒ}tt|d|dƒI d H }|S | 	d¡dkrr|dkrr| 
d	¡d
 }t d¡}||v ritt | |¡¡ƒ}tt|d|dƒI d H }|S tdd|› d‚| d¡r·|dkr·t |d ¡}	|	 d¡}
|	 d¡}| d¡}|
|ks—J ‚t|
ƒ}|t| d¡< tt |¡ƒ}tt|d|dƒI d H }|S | d¡rè|dkrèg }t ¡ D ]}t |¡}| |¡ qÆtt |¡ƒ}tt|d|dƒI d H }|S | 	d¡dkrU|dkrU| 
d¡d
  
d	¡\}}}|dkr!|tv rtt|dƒƒI d H }|S tdd|› d‚|dkrL|tv rCtt t |¡¡ƒ}tt|d|dƒI d H }|S tdd|› d‚tdd|› d‚| 	d¡dkrz|dkrz| 
d	¡d
 }t |¡ tt|dƒƒI d H }|S | 	d¡dkr´|dkr¶| 
d	¡d
 }|tv r«tt t |¡¡ƒ}tt|d|dƒI d H }|S tdd|› d‚d S d S )NÚGETÚmethodÚurlú/api/kernelspecséÈ   )Úbufferz/api/kernelspecs/r   ú/é   r!   é”  zKernelspec does not exist: )Úmessageú/api/kernelsÚPOSTÚbodyr   r   ÚKERNEL_KSPEC_NAMEr$   éÉ   ú/api/kernels/Ú	interruptéÌ   zKernel does not exist: ÚrestartzBad action detected: ÚDELETEr   )r   r*   Úendswithr   ÚjsonÚdumpsr!   r	   r   ÚrfindÚ
rpartitionÚgetr   Úloadsr0   Úrunning_kernelsÚkeysÚappendÚpop)r3   Úkwargsr2   ÚrequestÚendpointÚresponse_bufÚresponseÚrequested_kernelspecÚkspecsÚ	json_bodyr   r   Ú
kspec_namer/   Úkernelsr.   Úrequested_kernel_idÚsepÚactionr   r   r   Úmock_gateway_request(   s€   €










ùr]   z)notebook.gateway.managers.gateway_requestc                       sÌ   e Zd ZdZdZe‡ fdd„ƒZe‡ fdd„ƒZe‡ fdd„ƒZe‡ fd	d
„ƒZ	‡ fdd„Z
dd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd „ Zd!d"„ Zd#d$„ Zd%d&„ Z‡  ZS )'ÚTestGatewayzhttp://mock-gateway-server:8889Úalicec                    ó   t  ¡  tƒ  ¡  d S ©N)r   Úclear_instanceÚsuperÚsetup_class©Úcls©Ú	__class__r   r   rd   ‡   ó   zTestGateway.setup_classc                    r`   ra   )r   rb   rc   Úteardown_classre   rg   r   r   rj   Œ   ri   zTestGateway.teardown_classc                    s    t ƒ  ¡ }| tjddœ¡ |S )Nz44.4)ÚJUPYTER_GATEWAY_URLÚJUPYTER_GATEWAY_CONNECT_TIMEOUT)rc   Úget_patch_envÚupdater^   Úmock_gateway_url)rf   Útest_envrg   r   r   rm   ‘   s
   
ÿzTestGateway.get_patch_envc                    s"   t ƒ  ¡ }| ddtj g¡ |S )Nz$--GatewayClient.request_timeout=96.0z--GatewayClient.http_user=)rc   Úget_argvÚextendr^   Úmock_http_user)rf   r   rg   r   r   rq   ˜   s   
zTestGateway.get_argvc                    s(   t ƒ }t ¡ jdi |¤Ž tƒ  ¡  d S )Nr   )Údictr   ÚinstanceÚload_connection_argsrc   ÚsetUp)ÚselfrP   rg   r   r   rw   ž   s   zTestGateway.setUpc                 C   s   | j jjdks	J ‚| j jjtjksJ ‚| j jjtjksJ ‚| j jj| j jjks)J ‚| j jjdks2J ‚| j jj	dks;J ‚t
jd tdƒksFJ ‚d S )NTg333333F@g      X@ÚKERNEL_LAUNCH_TIMEOUTé`   )ÚnotebookÚgateway_configÚgateway_enabledr3   r^   ro   Ú	http_userrs   Úconnect_timeoutÚrequest_timeoutÚosÚenvironr*   ©rx   r   r   r   Útest_gateway_options£   s   z TestGateway.test_gateway_optionsc                 C   s:   | j jjdks	J ‚| j jjdksJ ‚| j jjdksJ ‚d S )NÚGatewayKernelManagerÚGatewaySessionManagerÚGatewayKernelSpecManager)r{   Úkernel_manager_classÚ__name__Úsession_manager_classÚkernel_spec_manager_classrƒ   r   r   r   Útest_gateway_class_mappings¬   s   z'TestGateway.test_gateway_class_mappingsc                 C   sˆ   t 8 |  dd¡}|  |jd¡ t |j d¡¡}| d¡}t	|ƒdks&J ‚| d¡ d¡dks2J ‚W d   ƒ d S 1 s=w   Y  d S )	Nr1   r4   r5   úutf-8r!   r8   r   r   )
Úmocked_gatewayrQ   ÚassertEqualÚstatus_coderF   rK   ÚcontentÚdecoderJ   Úlen)rx   rT   r‘   rV   r   r   r   Útest_gateway_get_kernelspecs²   s   
"úz(TestGateway.test_gateway_get_kernelspecsc                 C   s‚   t 5 |  dd¡}|jdksJ ‚t |j d¡¡}| d¡dks"J ‚|  dd¡}|jdks/J ‚W d   ƒ d S 1 s:w   Y  d S )	Nr1   z/api/kernelspecs/kspec_foor5   r   r   r   z/api/kernelspecs/no_such_specr9   )rŽ   rQ   r   rF   rK   r‘   r’   rJ   )rx   rT   r   r   r   r   Ú!test_gateway_get_named_kernelspec¼   s   "ùz-TestGateway.test_gateway_get_named_kernelspecc                 C   sp   |   d¡\}}|  |  |¡¡ |  |¡ |  |  |¡¡ |  |¡ |  |  |¡¡ |  |¡ |  |  |¡¡ d S )Nr   )Úcreate_sessionÚ
assertTrueÚis_kernel_runningÚinterrupt_kernelÚrestart_kernelÚdelete_sessionÚassertFalse)rx   Ú
session_idr.   r   r   r   Útest_gateway_session_lifecycleÇ   s   


z*TestGateway.test_gateway_session_lifecyclec                 C   sl   |   d¡}|  |  |¡¡ |  |¡ |  |  |¡¡ |  |¡ |  |  |¡¡ |  |¡ |  |  |¡¡ d S )Nr   )Úcreate_kernelr—   r˜   r™   rš   Údelete_kernelrœ   )rx   r.   r   r   r   Útest_gateway_kernel_lifecycleà   s   



z)TestGateway.test_gateway_kernel_lifecyclec           	      C   s  t w tj | jd¡}tƒ }|dd|idœ|d< |tjd< | jdi |¤Ž}|  |j	d	¡ t
 |j d
¡¡}|  | d¡|¡ | d¡ d¡}t |¡}|  || d¡¡ |  | d¡ d¡| d¡¡ | d¡}tj d¡ ||fW  d  ƒ S 1 s|w   Y  dS )zšCreates a session for a kernel.  The session is created against the notebook server
           which then uses the gateway for kernel management.
        ztestgw.ipynbr{   r   )ÚpathÚtypeÚkernelrF   r>   r<   ú/api/sessionsr?   r   r¢   r¤   r$   N)r<   r¥   )rŽ   r   r¢   ÚjoinÚnotebook_dirrt   r‚   rQ   r   r   rF   rK   r‘   r’   rJ   rL   rO   )	rx   Úkernel_nameÚnb_pathrP   rT   r/   r.   Úrunning_kernelr   r   r   r   r–   ù   s"   


$êzTestGateway.create_sessionc                 C   óV   t  |  dd| ¡}|  |jd¡ |  |jd¡ W d  ƒ dS 1 s$w   Y  dS )zADeletes a session corresponding to the given session id.
        rD   z/api/sessions/rB   ú
No ContentN©rŽ   rQ   r   r   Úreason)rx   r   rT   r   r   r   r›     ó
   "üzTestGateway.delete_sessionc                 C   s–   t ? |  dd¡}|  |jd¡ t |j d¡¡}|  t|ƒtt	ƒ¡ |D ]}| 
d¡|kr7 W d  ƒ dS q%	 W d  ƒ dS 1 sDw   Y  dS )	z9Issues request to get the set of running kernels
        r1   r;   r5   r   r$   NTF)rŽ   rQ   r   r   rF   rK   r‘   r’   r“   rL   rJ   )rx   r.   rT   rY   r/   r   r   r   r˜     s   ø$÷zTestGateway.is_kernel_runningc                 C   s¼   t R tƒ }d|i|d< |tjd< | jd
i |¤Ž}|  |jd¡ t |j	 
d¡¡}| d¡}t |¡}|  || d¡¡ |  | d¡|¡ tj d¡ |W  d	  ƒ S 1 sWw   Y  d	S )ú3Issues request to restart the given kernel
        r   rF   r>   r<   r;   r?   r   r$   N)r<   r;   )rŽ   rt   r   r‚   rQ   r   r   rF   rK   r‘   r’   rJ   rL   rO   )rx   r¨   rP   rT   r/   r.   rª   r   r   r   rŸ   ,  s   


$îzTestGateway.create_kernelc                 C   sZ   t ! |  dd| d ¡}|  |jd¡ |  |jd¡ W d  ƒ dS 1 s&w   Y  dS )z5Issues request to interrupt the given kernel
        r<   r@   z
/interruptrB   r¬   Nr­   ©rx   r.   rT   r   r   r   r™   C  s
   "ýzTestGateway.interrupt_kernelc                 C   sœ   t B |  dd| d ¡}|  |jd¡ t |j d¡¡}| d¡}t	 |¡}|  || d¡¡ |  | d¡| d¡¡ W d  ƒ dS 1 sGw   Y  dS )	r°   r<   r@   z/restartr5   r   r$   r   N)
rŽ   rQ   r   r   rF   rK   r‘   r’   rJ   rL   )rx   r.   rT   r/   Úrestarted_kernel_idrª   r   r   r   rš   K  s   

"øzTestGateway.restart_kernelc                 C   r«   )z=Deletes kernel corresponding to the given kernel id.
        rD   r@   rB   r¬   Nr­   r±   r   r   r   r    X  r¯   zTestGateway.delete_kernel)r‰   Ú
__module__Ú__qualname__ro   rs   Úclassmethodrd   rj   rm   rq   rw   r„   rŒ   r”   r•   rž   r¡   r–   r›   r˜   rŸ   r™   rš   r    Ú__classcell__r   r   rg   r   r^   ‚   s2    	
	r^   )Ú__doc__r   rF   r+   r   Úior   Úunittest.mockr   Útornado.webr   Útornado.httpclientr   r   Únotebook.gateway.managersr   Únotebook.utilsr	   Úlaunchnotebookr   r   r!   rt   rL   r0   r]   rŽ   r^   r   r   r   r   Ú<module>   s&    
W