
    uki"                         d 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	 ddl
mZ ddl
mZ ej                  Zej                  Z G d	 d
ej                        Zedk(  r ej$                          yy)z@Tests for mediapipe.python._framework_bindings.calculator_graph.    )absltest)text_format)calculator_pb2)packet_creator)packet_getter)calculator_graphvalidated_graph_configc                   B    e Zd Zd Zd Zd Zd Zd Zd Zd Z	d Z
d	 Zy
)	GraphTestc                 r    | j                  t        d      5  t        d       d d d        y # 1 sw Y   y xY w)Nz3(No such file or directory|The path does not exist)z/tmp/abc.binarypb)binary_graph_path)assertRaisesRegexFileNotFoundErrorCalculatorGraph)selfs    a/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/mediapipe/python/calculator_graph_test.pytest_invalid_binary_graph_filez(GraphTest.test_invalid_binary_graph_file    s6    			=
? = (;<= = =s   -6c                     d}t        j                         }t        j                  ||       | j	                  t
        d      5  t        |       d d d        y # 1 sw Y   y xY w)Nz
      node {
        calculator: 'PassThroughCalculator'
        input_stream: 'in'
        input_stream: 'in'
        output_stream: 'out'
      }
    zUInput and output streams to PassThroughCalculator must use matching tags and indexes.graph_config)r   CalculatorGraphConfigr   Parser   
ValueErrorr   r   text_configconfig_protos      r   test_invalid_node_configz"GraphTest.test_invalid_node_config&   s\    K "779Lk<0			_
 1 <0	1 1 1   AA"c                     d}t        j                         }t        j                  ||       | j	                  t
        d      5  t        |       d d d        y # 1 sw Y   y xY w)Nz~
      node {
        calculator: 'SomeUnknownCalculator'
        input_stream: 'in'
        output_stream: 'out'
      }
    z1Unable to find Calculator "SomeUnknownCalculator"r   )r   r   r   r   r   RuntimeErrorr   r   s      r   test_invalid_calculator_typez&GraphTest.test_invalid_calculator_type7   s[    K "779Lk<0			K
M 1<01 1 1r   c                    d}t        j                         }t        j                  ||       t	        |      }t        j                  d      }g t	        |      }|j                  dfd       |j                          |j                  d|d       |j                  d|j                  d	      
       |j                          | j                  |j                  t        j                  j                          | j                  |j"                  d	       | j%                  |j'                                | j)                  d       | j                  d   j*                  d       | j                  d	   j*                  d	       | j                  t-        j.                  d         d       | j                  t-        j.                  d	         d       y )N
      max_queue_size: 1
      input_stream: 'in'
      output_stream: 'out'
      node {
        calculator: 'PassThroughCalculator'
        input_stream: 'in'
        output_stream: 'out'
      }
    r   hello worldoutc                 &    j                  |      S Nappend_packetr&   s     r   <lambda>zDGraphTest.test_graph_initialized_with_proto_config.<locals>.<lambda>W       F9K     inr   streamr-   	timestamp   r3   r-      )r   r   r   r   r   r   create_stringobserve_output_stream	start_runadd_packet_to_input_streamatcloseassertEqualgraph_input_stream_add_moder   GraphInputStreamAddModeWAIT_TILL_NOT_FULLmax_queue_sizeassertFalse	has_error	assertLenr4   r   get_str)r   r   r   graphhello_world_packetr&   s        @r   (test_graph_initialized_with_proto_configz2GraphTest.test_graph_initialized_with_proto_configE   s   	K "779Lk<06E'55mD
C6E	'KL	OO	$$.! % =	$$.11!4 % 6	KKM))00CCE 	U))1-U__&'NN3SV%%q)SV%%q)]**3q62MB]**3q62MBr0   c                 D   d}t        j                  d      }g t        |      }|j                  dfd       |j	                          |j                  d|j                  d             |j                  d|d	
       |j                          | j                  |j                  t        j                  j                         | j                  |j                  d	       | j                  |j                                | j!                  d       | j                  d   j"                  d       | j                  d	   j"                  d	       | j                  t%        j&                  d         d       | j                  t%        j&                  d	         d       y )Nr$   r%   r   r&   c                 &    j                  |      S r(   r)   r+   s     r   r.   zCGraphTest.test_graph_initialized_with_text_config.<locals>.<lambda>x   r/   r0   r1   r   r6   r5   r2   r7   )r   r8   r   r9   r:   r;   r<   r=   r>   r?   r   r@   rA   rB   rC   rD   rE   r4   r   rF   )r   r   rH   rG   r&   s       @r   'test_graph_initialized_with_text_configz1GraphTest.test_graph_initialized_with_text_configi   sc   	K (55mD
C5E	'KL	OO	$$.11!4 % 6	$$.! % =	KKM))00CCE 	U))1-U__&'NN3SV%%q)SV%%q)]**3q62MB]**3q62MBr0   c                    d}t        j                  d      }g t               }| j                  |j	                                |j                  |       | j                  |j	                                t        |      }|j                  dfd       |j                          |j                  d|j                  d      	       |j                  d|d
       |j                          | j                  |j                  t        j                   j"                         | j                  |j$                  d
       | j                  |j'                                | j)                  d       | j                  d   j*                  d       | j                  d
   j*                  d
       | j                  t-        j.                  d         d       | j                  t-        j.                  d
         d       y )Nr$   r%   r   r	   r&   c                 &    j                  |      S r(   r)   r+   s     r   r.   zDGraphTest.test_graph_validation_and_initialization.<locals>.<lambda>   r/   r0   r1   r   r6   r5   r2   r7   )r   r8   ValidatedGraphConfigrC   initialized
initialize
assertTruer   r9   r:   r;   r<   r=   r>   r?   r   r@   rA   rB   rD   rE   r4   r   rF   )r   r   rH   validated_graphrG   r&   s        @r   (test_graph_validation_and_initializationz2GraphTest.test_graph_validation_and_initialization   s   	K (55mD
C*,O_0023K8OOO//12?CE	'KL	OO	$$.11!4 % 6	$$.! % =	KKM))00CCE 	U))1-U__&'NN3SV%%q)SV%%q)]**3q62MB]**3q62MBr0   c                    d}t        j                         }t        j                  ||       t	        j
                  d      }g t        |      }|j                  dfd       |j                          |j                  d|j                  d             |j                          |j                  d|j                  d             | j                  t        d	      5  |j                          d d d        y # 1 sw Y   y xY w)
Nr$   r%   r   r&   c                 &    j                  |      S r(   r)   r+   s     r   r.   zCGraphTest.test_insert_packets_with_same_timestamp.<locals>.<lambda>   r/   r0   r1   r   r6   z7Current minimum expected timestamp is 1 but received 0.)r   r   r   r   r   r8   r   r9   r:   r;   r<   wait_until_idler   r   )r   r   r   rH   rG   r&   s        @r   'test_insert_packets_with_same_timestampz1GraphTest.test_insert_packets_with_same_timestamp   s    	K "779Lk<0'55mD
C6E	'KL	OO	$$.11!4 % 6		$$.11!4 % 6			M
O   s    C::Dc                    d}t        j                         }t        j                  ||       t	        |      }|j                  dt        j                  d      i       |j                          | j                  |j                                | j                  t        j                  |j                  d            d       y )Nz
      node {
        calculator: 'StringToUint64Calculator'
        input_side_packet: "string"
        output_side_packet: "number"
      }
    r   string42)input_side_packetsnumber*   )r   r   r   r   r   r:   r   r8   wait_until_donerC   rD   r>   r   get_uintget_output_side_packet)r   r   r   rG   s       r   test_side_packet_graphz GraphTest.test_side_packet_graph   s    K "779Lk<06E	OO$n&B&B4&HI  K	U__&'u;;HEFLr0   c                    d}t        j                  d      }g t        |      }|j                  dfd       |j	                          d}t        |      D ]  }|j                  d||        |j                          | j                  |       t        |      D ]I  }| j                  |   j                  |       | j                  t        j                  |         d       K y )	Nr$   r%   r   r&   c                 &    j                  |      S r(   r)   r+   s     r   r.   z/GraphTest.test_sequence_input.<locals>.<lambda>   r/   r0   i  r1   r2   )r   r8   r   r9   r:   ranger;   rW   rE   r>   r4   r   rF   )r   r   rH   rG   sequence_sizeir&   s         @r   test_sequence_inputzGraphTest.test_sequence_input   s    	K (55mD
C5E	'KL	OOM=! ?&&0A ' ?? 
NN3&=! E
s1v''+
},,SV4mDEr0   N)__name__
__module____qualname__r   r   r"   rI   rL   rT   rX   rb   rh    r0   r   r   r      s7    =1"1"CHCB$CL8L$Er0   r   __main__N)__doc__absl.testingr   google.protobufr   mediapipe.frameworkr   mediapipe.pythonr   r   $mediapipe.python._framework_bindingsr   r
   r   rO   TestCaser   ri   mainrl   r0   r   <module>rv      sh    G ! ' . + * A G"22-BB YE!! YEx z(--/ r0   