# Path: android_env/environment.py
# class AndroidEnv(dm_env.Environment):
#   def __init__(self, coordinator: coordinator_lib.Coordinator):
#   def action_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def observation_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def task_extras_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def raw_action(self):
#   def raw_observation(self):
#   def android_logs(self) -> Dict[str, Any]:
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, action: Dict[str, np.ndarray]) -> dm_env.TimeStep:
#   def task_extras(self, latest_only: bool = True) -> Dict[str, np.ndarray]:
#   def close(self) -> None:
#   def __del__(self) -> None:
# 
# Path: android_env/components/coordinator.py
# class Coordinator():
#   def __init__(
#       self,
#       simulator: base_simulator.BaseSimulator,
#       task_manager: task_manager_lib.TaskManager,
#       step_timeout_sec: int = 10,
#       max_steps_per_sec: float = 15.0,
#       periodic_restart_time_min: float = 0.0,
#       force_simulator_launch: bool = True,
#   ):
#   def action_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def observation_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def task_extras_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def reset_environment_state(self):
#   def _lift_all_fingers(self) -> None:
#   def _should_periodic_restart(self) -> bool:
#   def _restart_simulator(self, max_retries: int = 3):
#   def execute_action(
#       self,
#       action: Optional[Dict[str, np.ndarray]],
#   ) -> Tuple[Optional[Dict[str, np.ndarray]], float, Dict[str, Any], bool]:
#   def _send_action_to_simulator(self, action: Dict[str, np.ndarray]) -> None:
#   def _check_timeout(self) -> bool:
#   def _wait_for_next_frame(self) -> None:
#   def _get_time_since_last_observation(self) -> float:
#   def get_logs(self) -> Dict[str, Any]:
#   def close(self):


from absl.testing import absltest
from android_env import environment
from android_env.components import coordinator as coordinator_lib
import dm_env
import mock
import numpy as np
    self.assertEqual(env.observation_spec()['timedelta'].shape, ())
    self.assertIn('orientation', env.observation_spec())
    # The `orientation` should be a one-hot vector with four dimensions.
    self.assertIsInstance(env.observation_spec()['orientation'],
                          dm_env.specs.Array)
    self.assertEqual(env.observation_spec()['orientation'].shape, (4,))

    # Check extras spec.
    self.assertNotEmpty(env.task_extras_spec())
    self.assertIn('click', env.task_extras_spec())
    self.assertEqual(env.task_extras_spec()['click'].shape, ())
    self.assertEqual(env.task_extras_spec()['click'].dtype, np.int64)

  def test_reset_and_step(self):
    coordinator = mock.create_autospec(coordinator_lib.Coordinator)
    coordinator.action_spec.return_value = {
        'action_type':
            dm_env.specs.DiscreteArray(num_values=3),
        'touch_position':
            dm_env.specs.BoundedArray(
                shape=(2,), dtype=np.float32, minimum=0.0, maximum=1.0),
    }
    coordinator.observation_spec.return_value = {
        'pixels': dm_env.specs.Array(shape=(123, 456, 3), dtype=np.uint8),
        'timedelta': dm_env.specs.Array(shape=(), dtype=np.int64),
        'orientation': dm_env.specs.Array(shape=(4,), dtype=np.uint8),
    }
    coordinator.task_extras_spec.return_value = {
        'click': dm_env.specs.Array(shape=(1,), dtype=np.int64),
    }

# Path: android_env/environment.py
# class AndroidEnv(dm_env.Environment):
#   def __init__(self, coordinator: coordinator_lib.Coordinator):
#   def action_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def observation_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def task_extras_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def raw_action(self):
#   def raw_observation(self):
#   def android_logs(self) -> Dict[str, Any]:
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, action: Dict[str, np.ndarray]) -> dm_env.TimeStep:
#   def task_extras(self, latest_only: bool = True) -> Dict[str, np.ndarray]:
#   def close(self) -> None:
#   def __del__(self) -> None:
# 
# Path: android_env/components/coordinator.py
# class Coordinator():
#   def __init__(
#       self,
#       simulator: base_simulator.BaseSimulator,
#       task_manager: task_manager_lib.TaskManager,
#       step_timeout_sec: int = 10,
#       max_steps_per_sec: float = 15.0,
#       periodic_restart_time_min: float = 0.0,
#       force_simulator_launch: bool = True,
#   ):
#   def action_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def observation_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def task_extras_spec(self) -> Dict[str, dm_env.specs.Array]:
#   def reset_environment_state(self):
#   def _lift_all_fingers(self) -> None:
#   def _should_periodic_restart(self) -> bool:
#   def _restart_simulator(self, max_retries: int = 3):
#   def execute_action(
#       self,
#       action: Optional[Dict[str, np.ndarray]],
#   ) -> Tuple[Optional[Dict[str, np.ndarray]], float, Dict[str, Any], bool]:
#   def _send_action_to_simulator(self, action: Dict[str, np.ndarray]) -> None:
#   def _check_timeout(self) -> bool:
#   def _wait_for_next_frame(self) -> None:
#   def _get_time_since_last_observation(self) -> float:
#   def get_logs(self) -> Dict[str, Any]:
#   def close(self):


# Path: android_env/environment_test.py
from absl.testing import absltest
from android_env import environment
from android_env.components import coordinator as coordinator_lib
import dm_env
import mock
import numpy as np

    self.assertEqual(env.observation_spec()['timedelta'].shape, ())
    self.assertIn('orientation', env.observation_spec())
    # The `orientation` should be a one-hot vector with four dimensions.
    self.assertIsInstance(env.observation_spec()['orientation'],
                          dm_env.specs.Array)
    self.assertEqual(env.observation_spec()['orientation'].shape, (4,))

    # Check extras spec.
    self.assertNotEmpty(env.task_extras_spec())
    self.assertIn('click', env.task_extras_spec())
    self.assertEqual(env.task_extras_spec()['click'].shape, ())
    self.assertEqual(env.task_extras_spec()['click'].dtype, np.int64)

  def test_reset_and_step(self):
    coordinator = mock.create_autospec(coordinator_lib.Coordinator)
    coordinator.action_spec.return_value = {
        'action_type':
            dm_env.specs.DiscreteArray(num_values=3),
        'touch_position':
            dm_env.specs.BoundedArray(
                shape=(2,), dtype=np.float32, minimum=0.0, maximum=1.0),
    }
    coordinator.observation_spec.return_value = {
        'pixels': dm_env.specs.Array(shape=(123, 456, 3), dtype=np.uint8),
        'timedelta': dm_env.specs.Array(shape=(), dtype=np.int64),
        'orientation': dm_env.specs.Array(shape=(4,), dtype=np.uint8),
    }
    coordinator.task_extras_spec.return_value = {
        'click': dm_env.specs.Array(shape=(1,), dtype=np.int64),
    }

    env = environment.AndroidEnv(coordinator)