
    Zi2                         d dl Z d dlZd dlZd dlmZmZmZ d dlZd dlm	Z	m
Z
mZmZ d dlmZmZmZmZ d dlmZmZ 	  G d dej                  Zedk    r ej                     dS dS )	    N)Mockpatch	MagicMock)recv	recv_linesendDEFAULT_SOCKET_OPTION)SSLErrorSSLEOFErrorSSLWantWriteErrorSSLWantReadError)WebSocketTimeoutException"WebSocketConnectionClosedExceptionc                       e Zd Zd Zd Zd Zd Zd Zd Zd Z	d Z
d	 Zd
 Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd Zd ZdS )
SocketTestc                     |                      t          t                     |                     t	          t                    d           t          d t          D                       }|                     |           dS )z4Test DEFAULT_SOCKET_OPTION contains expected optionsr   c              3   B   K   | ]}|d          t           j        k    V  dS )   N)socketTCP_NODELAY).0opts     V/root/projects/butler/venv/lib/python3.11/site-packages/websocket/tests/test_socket.py	<genexpr>z8SocketTest.test_default_socket_option.<locals>.<genexpr>.   s?        
  
-0CFf(( 
  
  
  
  
  
    N)assertIsInstancer	   listassertGreaterlenany
assertTrue)selftcp_nodelay_founds     r   test_default_socket_optionz%SocketTest.test_default_socket_option(   s    3T:::3455q999    
  
4I 
  
  
 
 
 	)*****r   c                     t                      }d|j        _        t          |d          }|                     |d           |j                            d           dS )zTest normal recv operation	   test data	   N)r   r   return_valueassertEqualassert_called_once_withr"   	mock_sockresults      r   test_recv_normalzSocketTest.test_recv_normal3   sV    FF	&2	#i##.....q11111r   c                 (   t                      }t          d          |j        _        |                     t
                    5 }t          |d           ddd           n# 1 swxY w Y   |                     t          |j                  d           dS )zTest recv with TimeoutErrorConnection timed outr'   N)	r   TimeoutErrorr   side_effectassertRaisesr   r)   str	exceptionr"   r,   cms      r   test_recv_timeout_errorz"SocketTest.test_recv_timeout_error=   s    FF	%12H%I%I	"899 	RA	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	R\**,BCCCCCs   AA#&A#c                    t                      }t          j        d          }d|_        ||j        _        d|j        _        |                     t                    5 }t	          |d           ddd           n# 1 swxY w Y   ddl
}|j        dk    r*|                     t          |j                  d           dS |                     t          |j                  d           dS )	zTest recv with socket.timeoutSocket timed out)r:         >@r'   Nr   )   
   r0   )r   r   timeoutargsr   r2   
gettimeoutr(   r3   r   sysversion_infor)   r4   r5   )r"   r,   timeout_excr7   rA   s        r   test_recv_socket_timeoutz#SocketTest.test_recv_socket_timeoutG   s   FF	n%7880%0	",0	)899 	RA	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	


w&&S..0FGGGGGS..0BCCCCCs   A99A= A=c                 :   t                      }t          d          }d|_        ||j        _        |                     t                    5 }t          |d           ddd           n# 1 swxY w Y   |                     t          |j	                  d           dS )z Test recv with SSL timeout errorThe operation timed out)rF   r'   N)
r   r
   r?   r   r2   r3   r   r)   r4   r5   )r"   r,   ssl_excr7   s       r   test_recv_ssl_timeoutz SocketTest.test_recv_ssl_timeout]   s    FF	4553%,	"899 	RA	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	R\**,EFFFFFs   A((A,/A,c                     t                      }t          d          }d|_        ||j        _        |                     t                    5  t          |d           ddd           dS # 1 swxY w Y   dS )z$Test recv with SSL non-timeout errorSSL certificate error)rJ   r'   N)r   r
   r?   r   r2   r3   )r"   r,   rG   s      r   test_recv_ssl_non_timeout_errorz*SocketTest.test_recv_ssl_non_timeout_errori   s    FF	2331%,	" x(( 	 	A	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	s   A))A-0A-c                    t                      }d|j        _        |                     t                    5 }t          |d           ddd           n# 1 swxY w Y   |                     t          |j                  d           dS )z1Test recv with empty response (connection closed)r   r'   Nz#Connection to remote host was lost.)r   r   r(   r3   r   r)   r4   r5   r6   s      r   test_recv_empty_responsez#SocketTest.test_recv_empty_responset   s    FF	&)	#ABB 	bA	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	R\**,QRRRRRs   AAAc                    t                      }t                      dg|j        _        t	          d          5 }t                      }||_        dg|j        _        t          |d          }|                     |d           |j        	                                 |j
        	                                 ddd           dS # 1 swxY w Y   dS )z.Test recv with SSLWantReadError (should retry)s   data after retryselectors.DefaultSelectorTd   N)r   r   r   r2   r   r(   selectr)   registerassert_calledcloser"   r,   mock_selector_classmock_selectorr-   s        r   test_recv_ssl_want_read_errorz(SocketTest.test_recv_ssl_want_read_error~   s   FF	 '7&8&8:M%N	".// 		03F FFM/<,15M -)S))FV%8999"00222--///		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0s   A;C  CCc                    t                      }t                      |j        _        d|j        _        t          d          5 }t                      }||_        g |j        _        |                     t                    5  t          |d           ddd           n# 1 swxY w Y   ddd           dS # 1 swxY w Y   dS )z.Test recv with SSLWantReadError that times outg      ?rO   rP   N)
r   r   r   r2   r@   r(   r   rQ   r3   r   )r"   r,   rV   rW   s       r   test_recv_ssl_want_read_timeoutz*SocketTest.test_recv_ssl_want_read_timeout   s.   FF	%5%7%7	",/	).// 	%3F FFM/<,02M -""#<== % %Y$$$% % % % % % % % % % % % % % %	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	% 	%s6   <B3>BB3B	B3"B	#B33B7:B7c                     t                      }g d}t          d|          5 }t          |          }|                     |d           |                     |j        d           ddd           dS # 1 swxY w Y   dS )zTest recv_line functionality)   H   e   lr^      o   
zwebsocket._socket.recv)r2   s   Hello
   N)r   r   r   r)   
call_count)r"   r,   
recv_calls	mock_recvr-   s        r   test_recv_linezSocketTest.test_recv_line   s    FF	 ;::
+DDD 	6	y))FVZ000Y11555		6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6 	6s   AA22A69A6c                     t                      }d|j        _        d|j        _        t          |d          }|                     |d           |j                            d           dS )zTest normal send operationr'   r;   r&   Nr   r   r(   r@   r)   assert_called_withr+   s      r   test_send_normalzSocketTest.test_send_normal   sa    FF	&'	#,0	)i..###)),77777r   c                     t                      }d|j        _        d|j        _        t          |d          }|                     |d           |j                            d           dS )z*Test send with zero timeout (non-blocking)r'   r   r&   N)r   r   r(   r@   r)   r*   r+   s      r   test_send_zero_timeoutz!SocketTest.test_send_zero_timeout   sa    FF	&'	#,-	)i..###..|<<<<<r   c                 @   t                      }d|j        _        t          d          |j        _        |                     t                    5 }t	          |d           ddd           n# 1 swxY w Y   |                     t          |j
                  d           dS )zTest send with SSLEOFErrorr;   zConnection closedr&   Nzsocket is already closed.)r   r@   r(   r   r   r2   r3   r   r)   r4   r5   r6   s      r   test_send_ssl_eof_errorz"SocketTest.test_send_ssl_eof_error   s    FF	,0	)%01D%E%E	"ABB 	*bL)))	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	R\**,GHHHHHs   A++A/2A/c                    t                      }d|j        _        t                      dg|j        _        t          d          5 }t                      }||_        dg|j        _        t	          |d          }|                     |d           |j	        
                                 |j        
                                 ddd           dS # 1 swxY w Y   dS )z/Test send with SSLWantWriteError (should retry)r;   r'   rO   Tr&   N)r   r@   r(   r   r   r2   r   rQ   r)   rR   rS   rT   rU   s        r   test_send_ssl_want_write_errorz)SocketTest.test_send_ssl_want_write_error   s   FF	,0	) '8&9&91%=	".// 		03F FFM/<,15M -)\22FVQ'''"00222--///		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0 		0s   A;CCCc                    t                      }d|j        _        t          j        d          }t
          j        |_        t
          j        df|_        |dg|j        _	        t          d          5 }t                      }||_        dg|j        _        t          |d          }|                     |d           ddd           dS # 1 swxY w Y   dS )z*Test send with EAGAIN error (should retry)r;    Resource temporarily unavailabler'   rO   Tr&   Nr   r@   r(   r   errorerrnoEAGAINr?   r   r2   r   rQ   r)   r"   r,   eagain_errorrV   rW   r-   s         r   test_send_socket_eagain_errorz(SocketTest.test_send_socket_eagain_error   s   FF	,0	) |$FGG"\"\+MN '3A%6	".// 	(3F FFM/<,15M -)\22FVQ'''	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	(   0A	CC
C
c                    t                      }d|j        _        t          j        d          }t
          j        |_        t
          j        df|_        |dg|j        _	        t          d          5 }t                      }||_        dg|j        _        t          |d          }|                     |d           ddd           dS # 1 swxY w Y   dS )z/Test send with EWOULDBLOCK error (should retry)r;   zOperation would blockr'   rO   Tr&   N)r   r@   r(   r   rs   rt   EWOULDBLOCKr?   r   r2   r   rQ   r)   )r"   r,   ewouldblock_errorrV   rW   r-   s         r   "test_send_socket_ewouldblock_errorz-SocketTest.test_send_socket_ewouldblock_error   s   FF	,0	) #L)@AA"'"3"'"35L!M '8%;	".// 	(3F FFM/<,15M -)\22FVQ'''	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	(ry   c                 R   t                      }d|j        _        t          j        d          }t
          j        |_        t
          j        df|_        ||j        _	        | 
                    t          j                  5  t          |d           ddd           dS # 1 swxY w Y   dS )z0Test send with other socket error (should raise)r;   zConnection reset by peerr&   N)r   r@   r(   r   rs   rt   
ECONNRESETr?   r   r2   r3   )r"   r,   other_errors      r   test_send_socket_other_errorz'SocketTest.test_send_socket_other_error  s    FF	,0	) l#=>>!,!,.HI%0	"v|,, 	* 	*L)))	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	*s   >BB #B c                    t                      }d|j        _        t          j        d          }d|_        ||j        _        |                     t          j                  5  t          |d           ddd           dS # 1 swxY w Y   dS )z-Test send with socket error that has no errnor;   Generic socket error)r   r&   N)	r   r@   r(   r   rs   r?   r   r2   r3   )r"   r,   no_errno_errors      r   test_send_socket_error_no_errnoz*SocketTest.test_send_socket_error_no_errno  s    FF	,0	)  &<==7%3	"v|,, 	* 	*L)))	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	* 	*s   !A??BBc                    t                      }d|j        _        t          j        d          }t
          j        |_        t
          j        df|_        ||j        _	        t          d          5 }t                      }||_        g |j        _        t          |d          }|                     |d           ddd           dS # 1 swxY w Y   dS )z!Test send write operation timeoutr;   rq   rO   r&   r   Nrr   rv   s         r   test_send_write_timeoutz"SocketTest.test_send_write_timeout(  s   FF	,0	) |$FGG"\"\+MN%1	".// 	(3F FFM/<,02M -)\22F VQ'''	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	( 	(s   .ACC
Cc                     t                      }d|j        _        d|j        _        t          |d          }|                     |d           |j                            d           dS )z.Test send with string data (should be encoded)r'   r;   z	test datar&   Nrg   r+   s      r   test_send_string_dataz SocketTest.test_send_string_data>  sa    FF	&'	#,0	)i--###)),77777r   c                 <   t                      }d|j        _        t          j        d          }t
          j        |_        t
          j        df|_        |dg|j        _	        t          d          5 }t                      }||_        dg|j        _        t          |d          }|                     |d           |j                                         |j                                         |j                                         ddd           dS # 1 swxY w Y   dS )zTest send retry mechanismr;   rq   r'   rO   Tr&   N)r   r@   r(   r   rs   rt   ru   r?   r   r2   r   rQ   r)   rR   rS   rT   rv   s         r   test_send_partial_send_retryz'SocketTest.test_send_partial_send_retryI  sG   FF	,0	) |$FGG"\"\+MN '3A%6	".// 	03F FFM/<,15M -)\22FVQ'''"00222 ..000--///	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0 	0s   0BDDDN)__name__
__module____qualname__r$   r.   r8   rD   rH   rK   rM   rX   rZ   re   ri   rk   rm   ro   rx   r}   r   r   r   r   r    r   r   r   r   '   s^       	+ 	+ 	+2 2 2D D DD D D,
G 
G 
G	 	 	S S S0 0 0$% % %6 6 6	8 	8 	8	= 	= 	=	I 	I 	I0 0 0&( ( (,( ( (,* * ** * *( ( (,	8 	8 	80 0 0 0 0r   r   __main__)rt   r   unittestunittest.mockr   r   r   timewebsocket._socketr   r   r   r	   websocket._ssl_compatr
   r   r   r   websocket._exceptionsr   r   TestCaser   r   mainr   r   r   <module>r      s0      0 0 0 0 0 0 0 0 0 0  J J J J J J J J J J J J                  
&z0 z0 z0 z0 z0" z0 z0 z0z	 zHMOOOOO r   