Merge pull request #206 from inverted-ai/update-mock-api

Update mock api

invertedai/api/drive.py

@@ -9,6 +9,7 @@
     mock_update_agent_state,
     get_mock_birdview,
     get_mock_infractions,
+    get_mock_light_recurrent_states
 )
 from invertedai.error import APIConnectionError, InvalidInput
 from invertedai.common import (
@@ -140,6 +141,9 @@ def drive(
             recurrent_states=recurrent_states,
             birdview=birdview,
             infractions=infractions,
+            traffic_lights_states=traffic_lights_states if traffic_lights_states is not None else None,
+            light_recurrent_states=get_mock_light_recurrent_states(len(traffic_lights_states)) if traffic_lights_states is not None else None,
+            api_model_version=api_model_version if api_model_version is not None else "best"
         )
         return response
 

invertedai/api/initialize.py

@@ -12,6 +12,8 @@
     get_mock_recurrent_state,
     get_mock_birdview,
     get_mock_infractions,
+    get_mock_traffic_light_states,
+    get_mock_light_recurrent_states
 )
 from invertedai.common import (
     RecurrentState,
@@ -128,12 +130,13 @@ def initialize(
     """
 
     if should_use_mock_api():
+        assert agent_attributes is not None or agent_count is not None
+        agent_count = agent_count if agent_count is not None else len(agent_attributes)
         if agent_attributes is None:
-            assert agent_count is not None
             agent_attributes = [get_mock_agent_attributes() for _ in range(agent_count)]
             agent_states = [get_mock_agent_state() for _ in range(agent_count)]
         else:
-            agent_states = states_history[-1]
+            agent_states = states_history[-1] if states_history is not None else []
         recurrent_states = [get_mock_recurrent_state() for _ in range(agent_count)]
         birdview = get_mock_birdview()
         infractions = get_mock_infractions(len(agent_states))
@@ -143,6 +146,9 @@ def initialize(
             recurrent_states=recurrent_states,
             birdview=birdview,
             infractions=infractions,
+            api_model_version=api_model_version if api_model_version is not None else "best",
+            traffic_lights_states=traffic_light_state_history[-1] if traffic_light_state_history is not None else None,
+            light_recurrent_states=get_mock_light_recurrent_states(len(traffic_light_state_history[0])) if traffic_light_state_history is not None else None
         )
         return response
 
@@ -160,7 +166,7 @@ def initialize(
         location_of_interest=location_of_interest,
         get_infractions=get_infractions,
         random_seed=random_seed,
-        model_version=api_model_version
+        model_version=api_model_version if api_model_version is not None else "best"
     )
     start = time.time()
     timeout = TIMEOUT

invertedai/api/location.py

@@ -66,6 +66,9 @@ def location_info(
             static_actors=[],
             bounding_polygon=[],
             max_agent_number=10,
+            map_center=Point(x=0, y=0),
+            map_fov=100,
+
         )
         return response