
    ߙi*                       d Z ddlmZ ddlZddlmc mZ ddl	Z	ddl
Z
ddlmZ ddlmZmZmZ ddlmZ ddlZ ee      j+                         j,                  d   Z ee      e
j2                  vr"e
j2                  j5                  d ee             ddlmZmZ dd	lmZm Z m!Z! dd
l"m#Z#m$Z$ ddl%m&Z&m'Z' ejP                  d        Z)ejP                  d        Z* G d d      Z+ejX                  j                   G d d             Z-ejX                  j                   G d d             Z.ejX                  j                   G d d             Z/y)zTests for State Machine and DAG Executor integration.

Tests for the integration between OrchestratorStateMachine and ActionExecutor
via the OrchestratorServer.
    )annotationsN)Path)	AsyncMock	MagicMockpatch)	dataclass   )OrchestratorServerServerConfig)OrchestratorStateOrchestratorStateMachineTransitionRecord)ActionExecutorExecutionReport)EdgeProxyClientClientConfigc                      t        dddd      S )zCreate a test server config.	localhostiA  FT)hostportenable_plannerenable_executor)r        [/home/nelsen/Projects/HRI/orchestrator/src/tests/test_state_machine_executor_integration.pyserver_configr       s     	 r   c                 8    ddddddddd	id
ddddddidgdgdS )zCreate a basic test plan.Navigate to lobbyDAGtest_plan_001step1	EXECUTIONrobotCHECK_BATTERY	min_level   step_idtype
agent_roleactionparamsstep2SPEAKmessageHello)r(   r)   r*   r+   r,   dependenciesintentworkflow_topologyplan_idstepsr   r   r   r   
basic_planr7   +   sW     &"" ##%)&+ ##%!$g.!(	
	 r   c                  (    e Zd ZdZd Zd Zd Zd Zy)$TestStateMachineListenerRegistrationz8Test that state machine listener is properly registered.c                   t               }|j                  }d}||k(  }|st        j                  d|fd||f      dt	        j
                         v st        j                  t               rt        j                  t               ndt        j                  |      t        j                  |      t        j                  |      dz  }dd|iz  }t        t        j                  |            dx}x}x}}y)	zDEdge proxy defaults should align with the robot-side server default.i  ==)zN%(py4)s
{%(py4)s = %(py2)s
{%(py2)s = %(py0)s()
}.edge_proxy_port
} == %(py7)sr   )py0py2py4py7assert %(py9)spy9N)
r   edge_proxy_port
@pytest_ar_call_reprcompare@py_builtinslocals_should_repr_global_name	_safereprAssertionError_format_explanation)self@py_assert1@py_assert3@py_assert6@py_assert5@py_format8@py_format10s          r   test_edge_proxy_default_portzATestStateMachineListenerRegistration.test_edge_proxy_default_portN   s    ~5~--55-5555-555555|555|555~555-5555555555r   c                   t        |      }|j                  }d}||u}|st        j                  d|fd||f      dt	        j
                         v st        j                  |      rt        j                  |      ndt        j                  |      t        j                  |      dz  }dd|iz  }t        t        j                  |            dx}x}}|j                  }d}t        ||      }	|	sd	d
t	        j
                         v st        j                  t              rt        j                  t              nd
dt	        j
                         v st        j                  |      rt        j                  |      ndt        j                  |      t        j                  |      t        j                  |	      dz  }t        t        j                  |            dx}x}}	|j                  }d}t        ||      }	|	sd	d
t	        j
                         v st        j                  t              rt        j                  t              nd
dt	        j
                         v st        j                  |      rt        j                  |      ndt        j                  |      t        j                  |      t        j                  |	      dz  }t        t        j                  |            dx}x}}	y)z;Test that executor is created when enable_executor is True.N)is not)z4%(py2)s
{%(py2)s = %(py0)s.executor
} is not %(py5)sserverr=   r>   py5assert %(py7)sr@   execute_planzSassert %(py7)s
{%(py7)s = %(py0)s(%(py3)s
{%(py3)s = %(py1)s.executor
}, %(py5)s)
}hasattr)r=   py1py3rX   r@   register_edge_proxy_client)r
   executorrD   rE   rF   rG   rH   rI   rJ   rK   r[   )
rL   r   rV   rM   @py_assert4rN   @py_format6rQ   @py_assert2rO   s
             r   "test_executor_created_when_enabledzGTestStateMachineListenerRegistration.test_executor_created_when_enabledR   s   #M2*d*d****d******v***v******d*******77w77777777w777w777777v777v7777777777777777E(DEw(DEEEEEEEEwEEEwEEEEEEvEEEvEEEEEE(DEEEEEEEEEEr   c                   t        d      }t        |      }|j                  }d}||u }|st        j                  d|fd||f      dt        j                         v st        j                  |      rt        j                  |      ndt        j                  |      t        j                  |      dz  }dd	|iz  }t        t        j                  |            dx}x}}y)
z@Test that executor is not created when enable_executor is False.F)r   Nis)z0%(py2)s
{%(py2)s = %(py0)s.executor
} is %(py5)srV   rW   rY   r@   )r   r
   r_   rD   rE   rF   rG   rH   rI   rJ   rK   )rL   configrV   rM   r`   rN   ra   rQ   s           r   'test_executor_not_created_when_disabledzLTestStateMachineListenerRegistration.test_executor_not_created_when_disabledZ   s    e4#F+&$&$&&&&$&&&&&&v&&&v&&&&&&$&&&&&&&r   c           	        t        |      }|j                  rV|j                  }|j                  }t	        |      }d}||kD  }|st        j                  d|fd||f      dt        j                         v st        j                  t              rt        j                  t              nddt        j                         v st        j                  |      rt        j                  |      ndt        j                  |      t        j                  |      t        j                  |      t        j                  |      dz  }dd|iz  }	t        t        j                  |	            d	x}x}x}x}}y	y	)
z=Test that state machine has listener registered for executor.r   )>)zs%(py7)s
{%(py7)s = %(py0)s(%(py5)s
{%(py5)s = %(py3)s
{%(py3)s = %(py1)s.state_machine
}._listeners
})
} > %(py10)slenrV   )r=   r\   r]   rX   r@   py10zassert %(py12)spy12N)r
   r_   state_machine
_listenersrk   rD   rE   rF   rG   rH   rI   rJ   rK   )
rL   r   rV   rb   r`   rO   @py_assert9@py_assert8@py_format11@py_format13s
             r   test_state_machine_has_listenerzDTestStateMachineListenerRegistration.test_state_machine_has_listener`   s    #M2??++;+66;367;!;7!;;;;7!;;;;;;3;;;3;;;;;;v;;;v;;;+;;;6;;;7;;;!;;;;;;;; r   N)__name__
__module____qualname____doc__rS   rc   rh   rt   r   r   r   r9   r9   K   s    B6F'<r   r9   c                      e Zd ZdZd Zd Zy) TestStateMachineTriggersExecutorz5Test that state machine transitions trigger executor.c                  K   t        |      }t        j                  |j                  _        ||_        |j                  j                  dd|i      }t        j                  }||k(  }|st        j                  d|fd||f      dt        j                         v st        j                  |      rt        j                  |      nddt        j                         v st        j                  t              rt        j                  t              ndt        j                  |      dz  }dd	|iz  }t        t        j                  |            d
x}}t!        j"                  d       d
{    y
7 w)z9Test that EXECUTING_ACTIONS transition triggers executor.user_confirmedplanr;   z9%(py0)s == %(py4)s
{%(py4)s = %(py2)s.EXECUTING_ACTIONS
}	new_stater   r=   r>   r?   assert %(py6)spy6Ng?)r
   r   AWAITING_CONFIRMrn   _state_pending_planhandle_eventEXECUTING_ACTIONSrD   rE   rF   rG   rH   rI   rJ   rK   asynciosleep)	rL   r   r7   rV   r   rN   rM   @py_format5@py_format7s	            r   (test_executing_actions_triggers_executorzITestStateMachineTriggersExecutor.test_executing_actions_triggers_executorq   s      $M2 '8&H&H#  * ((55vz2
	 .???y?????y???????y???y??????-???-??????????? mmC   s   E%E/'E-(E/c                :  
K   t        |      }g 
|j                  j                  
fd       t        j                  |j                  _        |j                  j                  |       d{   }|j                  r|j                  j                  d      }t        j                  }||k(  }|st        j                  d|fd||f      dt        j                         v st        j                  |      rt        j                   |      nddt        j                         v st        j                  t              rt        j                   t              ndt        j                   |      dz  }d	d
|iz  }	t#        t        j$                  |	            dx}}yy7 ,w)z:Test that executor finishing sends actions_complete event.c                &    j                  |       S Nappendttransitionss    r   <lambda>z`TestStateMachineTriggersExecutor.test_executor_finishes_sends_actions_complete.<locals>.<lambda>       K4F4Fq4I r   Nactions_completer;   z0%(py0)s == %(py4)s
{%(py4)s = %(py2)s.SPEAKING
}r   r   r   r   r   )r
   rn   add_listenerr   r   r   r_   rZ   successr   SPEAKINGrD   rE   rF   rG   rH   rI   rJ   rK   )rL   r   r7   rV   reportr   rN   rM   r   r   r   s             @r   -test_executor_finishes_sends_actions_completezNTestStateMachineTriggersExecutor.test_executor_finishes_sends_actions_complete   s     $M2 ))*IJ '8&I&I# 33J?? >>,,99:LMI 1 : ::9 :::::9 :::::::9:::9:::::: 1::: 1::: :::::::: 	 @s   A)F,F-D,FN)ru   rv   rw   rx   r   r   r   r   r   rz   rz   m   s    ?!,;r   rz   c                      e Zd ZdZd Zd Zy)TestEdgeProxyClientIntegrationz1Test Edge Proxy client integration with executor.c                <  K   t        |      }t               }t        d       |_        t        d       |_        t        d       |_        t        d       |_        |j                  r.|j                  j                  |       |j                  }|j                  }||u }|st        j                  d|fd||f      dt        j                         v st        j                  |      rt        j                  |      ndt        j                  |      t        j                  |      d	t        j                         v st        j                  |      rt        j                  |      nd	d
z  }dd|iz  }t        t        j                   |            dx}x}}yyw)z:Test that EdgeProxyClient can be registered with executor.c                    | S r   r   hs    r   r   zRTestEdgeProxyClientIntegration.test_edge_proxy_client_registered.<locals>.<lambda>       A r   side_effectc                    | S r   r   r   s    r   r   zRTestEdgeProxyClientIntegration.test_edge_proxy_client_registered.<locals>.<lambda>       Q r   c                    | S r   r   r   s    r   r   zRTestEdgeProxyClientIntegration.test_edge_proxy_client_registered.<locals>.<lambda>       q r   c                    | S r   r   r   s    r   r   zRTestEdgeProxyClientIntegration.test_edge_proxy_client_registered.<locals>.<lambda>       q r   re   )zR%(py4)s
{%(py4)s = %(py2)s
{%(py2)s = %(py0)s.executor
}._edge_client
} is %(py6)srV   mock_client)r=   r>   r?   r   assert %(py8)spy8N)r
   r   on_nav_statuson_robot_stateon_erroron_waypoint_listr_   r^   _edge_clientrD   rE   rF   rG   rH   rI   rJ   rK   )	rL   r   rV   r   rM   rN   rP   r   @py_format9s	            r   !test_edge_proxy_client_registeredz@TestEdgeProxyClientIntegration.test_edge_proxy_client_registered   s    #M2  k$-+$F!%.;%G"([A'0['I$ ??OO66{C??>?//>/;>>>>/;>>>>>>6>>>6>>>?>>>/>>>>>>;>>>;>>>>>>> s   FFc           
       K   t        |      }|j                  st        j                  d       t	               }t	        d       |_        t	        d       |_        t	        d       |_        t	        d       |_        t               |_
        |j                  j                  |       ddd	d
dddddddgd}t        j                  |j                  dd      5  |j                  j                  |       d{   }ddd       j                  }|syddt!        j"                         v st%        j&                  |      rt%        j(                  |      ndt%        j(                  |      dz  }t+        t%        j,                  |            d}|j                  }|j.                  }|sddt!        j"                         v st%        j&                  |      rt%        j(                  |      ndt%        j(                  |      t%        j(                  |      dz  }	t+        t%        j,                  |	            dx}}y7 B# 1 sw Y   BxY ww)z2Test executor with Edge Proxy client can navigate.Executor not enabledc                    | S r   r   r   s    r   r   zXTestEdgeProxyClientIntegration.test_executor_with_edge_proxy_navigates.<locals>.<lambda>   r   r   r   c                    | S r   r   r   s    r   r   zXTestEdgeProxyClientIntegration.test_executor_with_edge_proxy_navigates.<locals>.<lambda>   r   r   c                    | S r   r   r   s    r   r   zXTestEdgeProxyClientIntegration.test_executor_with_edge_proxy_navigates.<locals>.<lambda>   r   r   c                    | S r   r   r   s    r   r   zXTestEdgeProxyClientIntegration.test_executor_with_edge_proxy_navigates.<locals>.<lambda>   r   r   r   r   test_nav_planr!   r"   r#   NAVIGATElobbynormal)destinationspeedr'   r2   _wait_for_navigationN)return_value+assert %(py2)s
{%(py2)s = %(py0)s.success
}r   r=   r>   zVassert %(py4)s
{%(py4)s = %(py2)s
{%(py2)s = %(py0)s.send_navigate_waypoint
}.called
}r   r   )r
   r_   pytestskipr   r   r   r   r   r   send_navigate_waypointr^   r   objectrZ   r   rF   rG   rD   rH   rI   rJ   rK   called)
rL   r   rV   r   r}   r   rM   @py_format3rN   r   s
             r   'test_executor_with_edge_proxy_navigateszFTestEdgeProxyClientIntegration.test_executor_with_edge_proxy_navigates   s     $M2KK./  k$-+$F!%.;%G"([A'0['I$-6[*22;? *!&&  ''")(.5I	
  \\&//+APTU 	>!??77==F	> ~~~vv~1181888888888{888{88818888888888 >	> 	>s1   CII=H?>ID=I?IIIN)ru   rv   rw   rx   r   r   r   r   r   r   r      s    ;? '9r   r   c                      e Zd ZdZd Zd Zy)TestIntegrationz4Full integration tests for state machine + executor.c                j  K   t        |      }g |j                  j                  fd       t        j                  |j                  _        |j                  j                  dd|i      }t        j                  }||k(  }|st        j                  d|fd||f      dt        j                         v st        j                  |      rt        j                  |      nddt        j                         v st        j                  t              rt        j                  t              ndt        j                  |      dz  }d	d
|iz  }t        t        j                  |            dx}}|j                   j#                  |       d{   }	|	j$                  }|syddt        j                         v st        j                  |	      rt        j                  |	      ndt        j                  |      dz  }
t        t        j                  |
            d}|j                  j                  d      }t        j&                  }||k(  }|st        j                  d|fd||f      dt        j                         v st        j                  |      rt        j                  |      nddt        j                         v st        j                  t              rt        j                  t              ndt        j                  |      dz  }d	d
|iz  }t        t        j                  |            dx}}t)              }d}||k(  }|st        j                  d|fd||f      dt        j                         v st        j                  t(              rt        j                  t(              nddt        j                         v st        j                        rt        j                        ndt        j                  |      t        j                  |      dz  }dd|iz  }t        t        j                  |            dx}x}}d   }|j*                  }t        j                  }||k(  }|st        j                  d|fd||f      t        j                  |      t        j                  |      dt        j                         v st        j                  t              rt        j                  t              ndt        j                  |      dz  }dd|iz  }t        t        j                  |            dx}x}x}}d   }|j*                  }t        j&                  }||k(  }|st        j                  d|fd||f      t        j                  |      t        j                  |      dt        j                         v st        j                  t              rt        j                  t              ndt        j                  |      dz  }dd|iz  }t        t        j                  |            dx}x}x}}y7 w) zBTest full flow: AWAITING_CONFIRM -> EXECUTING_ACTIONS -> SPEAKING.c                &    j                  |       S r   r   r   s    r   r   zOTestIntegration.test_full_flow_confirm_to_execute_to_speaking.<locals>.<lambda>   r   r   r|   r}   r;   r~   state1r   r   r   r   Nr   r   r   r   r   state2r	   )z0%(py3)s
{%(py3)s = %(py0)s(%(py1)s)
} == %(py6)srk   r   )r=   r\   r]   r   r   r   r   )zS%(py3)s
{%(py3)s = %(py1)s.next
} == %(py7)s
{%(py7)s = %(py5)s.EXECUTING_ACTIONS
})r\   r]   rX   r@   rA   rB      )zJ%(py3)s
{%(py3)s = %(py1)s.next
} == %(py7)s
{%(py7)s = %(py5)s.SPEAKING
})r
   rn   r   r   r   r   r   r   rD   rE   rF   rG   rH   rI   rJ   rK   r_   rZ   r   r   rk   next)rL   r   r7   rV   r   rN   rM   r   r   r   r   r   rb   rP   r`   r   @py_assert0rO   rQ   rR   r   s                       @r   -test_full_flow_confirm_to_execute_to_speakingz=TestIntegration.test_full_flow_confirm_to_execute_to_speaking   s     $M2))*IJ '8&H&H# %%22vz2
 +<<<v<<<<<v<<<<<<<v<<<v<<<<<<*<<<*<<<<<<<<<<< 33J??~~~vv~ %%223EF*333v33333v3333333v333v333333*333*33333333333 ;$1$1$$$$1$$$$$$s$$$s$$$$$$;$$$;$$$$$$1$$$$$$$1~I~""I&7&I&II"&IIIII"&IIII~III"IIIIII&7III&7III&IIIIIIII1~@~""@&7&@&@@"&@@@@@"&@@@@~@@@"@@@@@@&7@@@&7@@@&@@@@@@@@ @s   FX3X0R)X3c                  
K   t        |      }|j                  st        j                  d       g 
|j                  j                  
fd       t        j                  |j                  _        ddddddd	d
gd}|j                  j                  |       d{   }|j                  s|j                  j                  d      }t        j                  }||k(  }|st        j                  d|fd||f      dt        j                          v st        j"                  |      rt        j$                  |      nddt        j                          v st        j"                  t              rt        j$                  t              ndt        j$                  |      dz  }dd|iz  }	t'        t        j(                  |	            dx}}yy7 ,w)z6Test that failed execution sends actions_failed event.r   c                &    j                  |       S r   r   r   s    r   r   zLTestIntegration.test_failed_execution_sends_actions_failed.<locals>.<lambda>  r   r   zInvalid actionr   invalid_planr!   r"   r#   INVALID_ACTION)r(   r)   r*   r+   r2   Nactions_failedr;   r   r   r   r   r   r   )r
   r_   r   r   rn   r   r   r   r   rZ   r   r   r   rD   rE   rF   rG   rH   rI   rJ   rK   )rL   r   rV   r   r   r   rN   rM   r   r   r   s             @r   *test_failed_execution_sends_actions_failedz:TestIntegration.test_failed_execution_sends_actions_failed	  sF     $M2KK./))*IJ&7&I&I# '!&%  ''").		
 33LAA ~~,,99:JKI 1 : ::9 :::::9 :::::::9:::9:::::: 1::: 1::: :::::::: 	 Bs   BG	GD,G	N)ru   rv   rw   rx   r   r   r   r   r   r   r      s    >A>#;r   r   )0rx   
__future__r   builtinsrF   _pytest.assertion.rewrite	assertionrewriterD   r   syspathlibr   unittest.mockr   r   r   dataclassesr   r   __file__resolveparentsROOTstrpathinsert
src.serverr
   r   src.state_machiner   r   r   src.executor.action_executorr   r   src.edge_proxy.clientr   r   fixturer   r7   r9   markrz   r   r   r   r   r   <module>r      s#   #    
  5 5 ! H~''*t9CHHHHOOAs4y! 7 [ [ H ?    >< <D -; -; -;j :9 :9 :9D E; E; E;r   