# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Tests for dm_env_rpc/dm_env adaptor."""

from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import dm_env
from dm_env import specs
import numpy as np

from google.rpc import status_pb2
from google.protobuf import text_format
from dm_env_rpc.v1 import dm_env_adaptor
from dm_env_rpc.v1 import dm_env_rpc_pb2
from dm_env_rpc.v1 import error
from dm_env_rpc.v1 import tensor_utils

_SAMPLE_STEP_REQUEST = dm_env_rpc_pb2.StepRequest(
    requested_observations=[1, 2],
    actions={
        1: tensor_utils.pack_tensor(4, dtype=dm_env_rpc_pb2.UINT8),
        2: tensor_utils.pack_tensor('hello')
    })
_SAMPLE_STEP_RESPONSE = dm_env_rpc_pb2.StepResponse(
    state=dm_env_rpc_pb2.EnvironmentStateType.RUNNING,
    observations={
        1: tensor_utils.pack_tensor(5, dtype=dm_env_rpc_pb2.UINT8),
        2: tensor_utils.pack_tensor('goodbye')
    })
_TERMINATED_STEP_RESPONSE = dm_env_rpc_pb2.StepResponse(
    state=dm_env_rpc_pb2.EnvironmentStateType.TERMINATED,
    observations={
        1: tensor_utils.pack_tensor(5, dtype=dm_env_rpc_pb2.UINT8),
        2: tensor_utils.pack_tensor('goodbye')
    })
_SAMPLE_SPEC = dm_env_rpc_pb2.ActionObservationSpecs(
    actions={
        1: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.UINT8, name='foo'),
        2: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.STRING, name='bar')
    },
    observations={
        1: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.UINT8, name='foo'),
        2: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.STRING, name='bar')
    },
)
_SAMPLE_SPEC_REORDERED = dm_env_rpc_pb2.ActionObservationSpecs(
    observations={
        2: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.STRING, name='bar'),
        1: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.UINT8, name='foo')
    },
    actions={
        2: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.STRING, name='bar'),
        1: dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.UINT8, name='foo')
    },
)
_SAMPLE_NESTED_SPECS = dm_env_rpc_pb2.ActionObservationSpecs(
    actions={
        1:
            dm_env_rpc_pb2.TensorSpec(
                dtype=dm_env_rpc_pb2.INT32, name='foo.bar'),
        2:
            dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.STRING, name='baz')
    },
    observations={
        1:
            dm_env_rpc_pb2.TensorSpec(
                dtype=dm_env_rpc_pb2.INT32, name='foo.bar'),
        2:
            dm_env_rpc_pb2.TensorSpec(dtype=dm_env_rpc_pb2.STRING, name='baz')
    },
)

# Ensures the equality check in reset() works if the dictionary elements are
# created in a different order.
_SAMPLE_RESET_RESPONSE = dm_env_rpc_pb2.ResetResponse(
    specs=_SAMPLE_SPEC_REORDERED)
_RESET_CHANGES_SPEC_RESPONSE = dm_env_rpc_pb2.ResetResponse(
    specs=dm_env_rpc_pb2.ActionObservationSpecs())

_RESERVED_SPEC = dm_env_rpc_pb2.ActionObservationSpecs(
    actions={},
    observations={
        1:
            dm_env_rpc_pb2.TensorSpec(
                dtype=dm_env_rpc_pb2.UINT8,
                name=dm_env_adaptor.DEFAULT_REWARD_KEY),
        2:
            dm_env_rpc_pb2.TensorSpec(
                dtype=dm_env_rpc_pb2.STRING,
                name=dm_env_adaptor.DEFAULT_DISCOUNT_KEY)
    })
_RESERVED_STEP_RESPONSE = dm_env_rpc_pb2.StepResponse(
    state=dm_env_rpc_pb2.EnvironmentStateType.RUNNING,
    observations={
        1: tensor_utils.pack_tensor(5, dtype=dm_env_rpc_pb2.UINT8),
        2: tensor_utils.pack_tensor('goodbye')
    })
_RESET_CHANGES_SPEC_ERROR = r'''Environment changed spec after reset.
before: "actions {
  key: 1
  value {
    name: "foo"
    dtype: UINT8
  }
}
actions {
  key: 2
  value {
    name: "bar"
    dtype: STRING
  }
}
observations {
  key: 1
  value {
    name: "foo"
    dtype: UINT8
  }
}
observations {
  key: 2
  value {
    name: "bar"
    dtype: STRING
  }
}
"
after: ""'''
_EXTENSIONS = {'baz': 123, 'qux': 'quux'}


class DmEnvAdaptorTests(absltest.TestCase):

  def setUp(self):
    super(DmEnvAdaptorTests, self).setUp()
    self._connection = mock.MagicMock()
    self._env = dm_env_adaptor.DmEnvAdaptor(
        self._connection, _SAMPLE_SPEC, extensions=_EXTENSIONS)

  def test_requested_observations(self):
    requested_observations = ['foo']
    filtered_env = dm_env_adaptor.DmEnvAdaptor(self._connection, _SAMPLE_SPEC,
                                               requested_observations)

    expected_filtered_step_request = dm_env_rpc_pb2.StepRequest(
        requested_observations=[1],
        actions={
            1: tensor_utils.pack_tensor(4, dtype=dm_env_rpc_pb2.UINT8),
            2: tensor_utils.pack_tensor('hello')
        })

    self._connection.send = mock.MagicMock(return_value=_SAMPLE_STEP_RESPONSE)
    filtered_env.step({'foo': 4, 'bar': 'hello'})

    self._connection.send.assert_called_once_with(
        expected_filtered_step_request)

  def test_invalid_requested_observations(self):
    requested_observations = ['invalid']
    with self.assertRaisesRegex(ValueError,
                                'Unsupported observations requested'):
      dm_env_adaptor.DmEnvAdaptor(self._connection, _SAMPLE_SPEC,
                                  requested_observations)

  def test_requested_observation_spec(self):
    requested_observations = ['foo']
    filtered_env = dm_env_adaptor.DmEnvAdaptor(self._connection, _SAMPLE_SPEC,
                                               requested_observations)
    observation_names = [name for name in filtered_env.observation_spec()]
    self.assertEqual(requested_observations, observation_names)

  def test_first_running_step(self):
    self._connection.send = mock.MagicMock(return_value=_SAMPLE_STEP_RESPONSE)
    timestep = self._env.step({'foo': 4, 'bar': 'hello'})

    self._connection.send.assert_called_once_with(_SAMPLE_STEP_REQUEST)
    self.assertEqual(dm_env.StepType.FIRST, timestep.step_type)
    self.assertIsNone(timestep.reward)
    self.assertIsNone(timestep.discount)
    self.assertEqual({'foo': 5, 'bar': 'goodbye'}, timestep.observation)

  def test_mid_running_step(self):
    self._connection.send = mock.MagicMock(return_value=_SAMPLE_STEP_RESPONSE)
    self._env.step({'foo': 4, 'bar': 'hello'})
    self._connection.send.assert_called_once_with(_SAMPLE_STEP_REQUEST)
    timestep = self._env.step({'foo': 4, 'bar': 'hello'})

    self.assertEqual(dm_env.StepType.MID, timestep.step_type)
    self.assertEqual(0.0, timestep.reward)
    self.assertEqual(1.0, timestep.discount)
    self.assertEqual({'foo': 5, 'bar': 'goodbye'}, timestep.observation)

  def test_last_step(self):
    self._connection.send = mock.MagicMock(return_value=_SAMPLE_STEP_RESPONSE)
    self._env.step({'foo': 4, 'bar': 'hello'})
    self._connection.send.assert_called_once_with(_SAMPLE_STEP_REQUEST)
    self._connection.send = mock.MagicMock(
        return_value=_TERMINATED_STEP_RESPONSE)
    timestep = self._env.step({'foo': 4, 'bar': 'hello'})

    self.assertEqual(dm_env.StepType.LAST, timestep.step_type)
    self.assertEqual(0.0, timestep.reward)
    self.assertEqual(0.0, timestep.discount)
    self.assertEqual({'foo': 5, 'bar': 'goodbye'}, timestep.observation)

  def test_illegal_state_transition(self):
    self._connection.send = mock.MagicMock(
        return_value=_TERMINATED_STEP_RESPONSE)
    with self.assertRaisesRegex(
        RuntimeError, 'Environment transitioned from TERMINATED to TERMINATED'):
      self._env.step({})

  def test_reset(self):
    self._connection.send = mock.MagicMock(return_value=_SAMPLE_STEP_RESPONSE)
    self._env.step({'foo': 4, 'bar': 'hello'})
    self._connection.send.assert_called_once_with(_SAMPLE_STEP_REQUEST)
    self._connection.send = mock.MagicMock(
        side_effect=[_SAMPLE_RESET_RESPONSE, _SAMPLE_STEP_RESPONSE])
    timestep = self._env.reset()

    self.assertEqual(dm_env.StepType.FIRST, timestep.step_type)
    self.assertIsNone(timestep.reward)
    self.assertIsNone(timestep.discount)
    self.assertEqual({'foo': 5, 'bar': 'goodbye'}, timestep.observation)

  def test_spec_generate_value_step(self):
    self._connection.send = mock.MagicMock(return_value=_SAMPLE_STEP_RESPONSE)
    action_spec = self._env.action_spec()
    actions = {
        name: spec.generate_value() for name, spec in action_spec.items()
    }
    self._env.step(actions)
    self._connection.send.assert_called_once_with(
        dm_env_rpc_pb2.StepRequest(
            requested_observations=[1, 2],
            actions={
                1: tensor_utils.pack_tensor(actions['foo']),
                2: tensor_utils.pack_tensor(actions['bar'], dtype=str)  # pytype: disable=wrong-arg-types  # typed-numpy
            }))

  def test_reset_changes_spec_raises_error(self):
    self._connection.send = mock.MagicMock(return_value=_SAMPLE_STEP_RESPONSE)
    self._env.step({'foo': 4, 'bar': 'hello'})
    self._connection.send.assert_called_once_with(_SAMPLE_STEP_REQUEST)
    self._connection.send = mock.MagicMock(
        side_effect=[_RESET_CHANGES_SPEC_RESPONSE, _SAMPLE_STEP_RESPONSE])
    with self.assertRaisesWithLiteralMatch(RuntimeError,
                                           _RESET_CHANGES_SPEC_ERROR):
      self._env.reset()

  def test_observation_spec(self):
    expected_spec = {
        'foo': specs.Array(shape=(), dtype=np.uint8, name='foo'),
        'bar': specs.StringArray(shape=(), name='bar')
    }
    self.assertEqual(expected_spec, self._env.observation_spec())

  def test_action_spec(self):
    expected_spec = {
        'foo': specs.Array(shape=(), dtype=np.uint8, name='foo'),
        'bar': specs.StringArray(shape=(), name='bar')
    }
    self.assertEqual(expected_spec, self._env.action_spec())

  def test_cant_step_after_close(self):
    self._connection.send = mock.MagicMock(
        return_value=dm_env_rpc_pb2.LeaveWorldResponse())
    self._env.close()
    with self.assertRaisesRegex(AttributeError, 'send'):
      self._env.step({})

  def test_reward_spec_default(self):
    self.assertEqual(
        specs.Array(shape=(), dtype=np.float64), self._env.reward_spec())

  def test_discount_spec_default(self):
    self.assertEqual(
        specs.BoundedArray(
            shape=(), dtype=np.float64, minimum=0.0, maximum=1.0),
        self._env.discount_spec())

  def test_close_leaves_world(self):
    self._connection.send = mock.MagicMock(
        return_value=dm_env_rpc_pb2.LeaveWorldResponse())
    self._env.close()
    self._connection.send.assert_called_once_with(
        dm_env_rpc_pb2.LeaveWorldRequest())

  def test_close_errors_when_cannot_leave_world(self):
    self._connection.send = mock.MagicMock(side_effect=ValueError('foo'))
    with self.assertRaisesRegex(ValueError, 'foo'):
      self._env.close()

  def test_close_releases_extensions(self):
    for extension_name, extension in _EXTENSIONS.items():
      self.assertEqual(getattr(self._env, extension_name), extension)
    self._connection.send = mock.MagicMock(
        return_value=dm_env_rpc_pb2.LeaveWorldResponse())
    self._env.close()
    for extension_name in _EXTENSIONS:
      self.assertIsNone(getattr(self._env, extension_name))

  def test_close_allows_multiple_close_calls(self):
    self._connection.send = mock.MagicMock(
        return_value=dm_env_rpc_pb2.LeaveWorldResponse(),
    )
    self._env.close()
    self._env.close()

    # Most importantly, when calling close multiple times, we need to make sure
    # that any _connection.send request is called once as the conenction is
    # destroyed after the first close.

    self._connection.send.assert_called_once_with(
        dm_env_rpc_pb2.LeaveWorldRequest(),
    )
    self.assertIsNone(self._env._connection)


class OverrideRewardDiscount(dm_env_adaptor.DmEnvAdaptor):

  def __init__(self):
    self.connection = mock.MagicMock()
    self.reward = mock.MagicMock()
    self.discount = mock.MagicMock()
    super(OverrideRewardDiscount, self).__init__(self.connection, _SAMPLE_SPEC)


class RewardDiscountOverrideTests(absltest.TestCase):

  def test_override_reward(self):
    env = OverrideRewardDiscount()
    env.reward.return_value = 0.5
    env.connection.send.return_value = _SAMPLE_STEP_RESPONSE
    timestep = env.step({})
    self.assertEqual(0.5, timestep.reward)
    env.reward.assert_called()
    self.assertEqual(dm_env_rpc_pb2.EnvironmentStateType.RUNNING,
                     env.reward.call_args[1]['state'])
    self.assertEqual(dm_env.StepType.FIRST,
                     env.reward.call_args[1]['step_type'])
    self.assertDictEqual({
        'foo': 5,
        'bar': 'goodbye'
    }, env.reward.call_args[1]['observations'])

  def test_override_discount(self):
    env = OverrideRewardDiscount()
    env.discount.return_value = 0.5
    env.connection.send.return_value = _SAMPLE_STEP_RESPONSE
    timestep = env.step({})
    self.assertEqual(0.5, timestep.discount)
    env.discount.assert_called()
    self.assertEqual(dm_env_rpc_pb2.EnvironmentStateType.RUNNING,
                     env.discount.call_args[1]['state'])
    self.assertEqual(dm_env.StepType.FIRST,
                     env.discount.call_args[1]['step_type'])
    self.assertDictEqual({
        'foo': 5,
        'bar': 'goodbye'
    }, env.discount.call_args[1]['observations'])


class ReservedKeywordTests(absltest.TestCase):

  def setUp(self):
    super(ReservedKeywordTests, self).setUp()
    self._connection = mock.MagicMock()
    self._env = dm_env_adaptor.DmEnvAdaptor(self._connection, _RESERVED_SPEC)

  def test_reward_spec(self):
    self.assertEqual(
        specs.Array(shape=(), dtype=np.uint8), self._env.reward_spec())

  def test_discount_spec(self):
    self.assertEqual(specs.StringArray(shape=()), self._env.discount_spec())

  def test_reward_from_reserved_keyword(self):
    self._connection.send = mock.MagicMock(return_value=_RESERVED_STEP_RESPONSE)
    self._env.step({})  # Reward is None for first step.
    timestep = self._env.step({})

    self.assertEqual(5, timestep.reward)

  def test_discount(self):
    self._connection.send = mock.MagicMock(return_value=_RESERVED_STEP_RESPONSE)
    timestep = self._env.step({})

    self.assertEqual('goodbye', timestep.discount)

  def test_observations_empty(self):
    self.assertEmpty(self._env.observation_spec())

  def test_explicitly_requesting_reward_and_discount(self):
    env = dm_env_adaptor.DmEnvAdaptor(
        self._connection,
        _RESERVED_SPEC,
        requested_observations=[
            dm_env_adaptor.DEFAULT_REWARD_KEY,
            dm_env_adaptor.DEFAULT_DISCOUNT_KEY
        ])
    expected_observation_spec = {
        dm_env_adaptor.DEFAULT_REWARD_KEY: env.reward_spec(),
        dm_env_adaptor.DEFAULT_DISCOUNT_KEY: env.discount_spec(),
    }
    self.assertEqual(env.observation_spec(), expected_observation_spec)


class EnvironmentAutomaticallyRequestsReservedKeywords(absltest.TestCase):

  def setUp(self):
    super(EnvironmentAutomaticallyRequestsReservedKeywords, self).setUp()
    self._connection = mock.MagicMock()
    self._env = dm_env_adaptor.DmEnvAdaptor(
        self._connection, _RESERVED_SPEC, requested_observations=[])
    self._connection.send = mock.MagicMock(return_value=_RESERVED_STEP_RESPONSE)

  def test_reward_spec_unrequested(self):
    self.assertEqual(
        specs.Array(shape=(), dtype=np.uint8), self._env.reward_spec())

  def test_discount_spec_unrequested(self):
    self.assertEqual(specs.StringArray(shape=()), self._env.discount_spec())

  def test_does_not_give_back_unrequested_observations(self):
    timestep = self._env.step({})
    self.assertEqual({}, timestep.observation)

  def test_first_reward_none(self):
    timestep = self._env.step({})
    self.assertIsNone(timestep.reward)

  def test_reward_piped_correctly(self):
    self._env.step({})  # Reward is None for first step.
    timestep = self._env.step({})
    self.assertEqual(5, timestep.reward)

  def test_discount_piped_correctly(self):
    timestep = self._env.step({})
    self.assertEqual('goodbye', timestep.discount)


class EnvironmentNestedActionsObservations(absltest.TestCase):

  def test_nested_specs(self):
    env = dm_env_adaptor.DmEnvAdaptor(
        connection=mock.MagicMock(), specs=_SAMPLE_NESTED_SPECS)
    expected_actions = {
        'foo': {
            'bar': specs.Array(shape=(), dtype=np.int32, name='foo.bar'),
        },
        'baz': specs.Array(shape=(), dtype=str, name='baz'),
    }
    expected_observations = {
        'foo': {
            'bar': specs.Array(shape=(), dtype=np.int32, name='foo.bar'),
        },
        'baz': specs.Array(shape=(), dtype=str, name='baz'),
    }

    self.assertSameElements(expected_actions, env.action_spec())
    self.assertSameElements(expected_observations, env.observation_spec())

  def test_no_nested_specs(self):
    env = dm_env_adaptor.DmEnvAdaptor(
        connection=mock.MagicMock(),
        specs=_SAMPLE_NESTED_SPECS,
        nested_tensors=False)
    expected_actions = {
        'foo.bar': specs.Array(shape=(), dtype=np.int32, name='foo.bar'),
        'baz': specs.Array(shape=(), dtype=str, name='baz'),
    }
    expected_observations = {
        'foo.bar': specs.Array(shape=(), dtype=np.int32, name='foo.bar'),
        'baz': specs.Array(shape=(), dtype=str, name='baz'),
    }

    self.assertSameElements(expected_actions, env.action_spec())
    self.assertSameElements(expected_observations, env.observation_spec())

  def test_nested_actions_step(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(
        return_value=text_format.Parse("""state: RUNNING""",
                                       dm_env_rpc_pb2.StepResponse()))
    env = dm_env_adaptor.DmEnvAdaptor(
        connection, specs=_SAMPLE_NESTED_SPECS, requested_observations=[])

    timestep = env.step({'foo': {'bar': 123}})

    self.assertEqual(dm_env.StepType.FIRST, timestep.step_type)

    connection.send.assert_called_once_with(
        text_format.Parse(
            """actions: { key: 1, value: { int32s: { array: 123 } } }""",
            dm_env_rpc_pb2.StepRequest()))

  def test_no_nested_actions_step(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(
        return_value=text_format.Parse("""state: RUNNING""",
                                       dm_env_rpc_pb2.StepResponse()))
    env = dm_env_adaptor.DmEnvAdaptor(
        connection,
        specs=_SAMPLE_NESTED_SPECS,
        requested_observations=[],
        nested_tensors=False)
    timestep = env.step({'foo.bar': 123})

    self.assertEqual(dm_env.StepType.FIRST, timestep.step_type)

    connection.send.assert_called_once_with(
        text_format.Parse(
            """actions: { key: 1, value: { int32s: { array: 123 } } }""",
            dm_env_rpc_pb2.StepRequest()))

  def test_nested_observations_step(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(
        return_value=text_format.Parse(
            """state: RUNNING
        observations: { key: 1, value: { int32s: { array: 42 } } }""",
            dm_env_rpc_pb2.StepResponse()))

    expected = {'foo': {'bar': 42}}

    env = dm_env_adaptor.DmEnvAdaptor(
        connection,
        specs=_SAMPLE_NESTED_SPECS,
        requested_observations=['foo.bar'])
    timestep = env.step({})
    self.assertEqual(dm_env.StepType.FIRST, timestep.step_type)
    self.assertSameElements(expected, timestep.observation)

    connection.send.assert_called_once_with(
        dm_env_rpc_pb2.StepRequest(requested_observations=[1]))

  def test_extensions(self):
    class _ExampleExtension:

      def foo(self):
        return 'bar'

    env = dm_env_adaptor.DmEnvAdaptor(
        connection=mock.MagicMock(),
        specs=_SAMPLE_SPEC,
        extensions={'extension': _ExampleExtension()})

    self.assertEqual('bar', env.extension.foo())

  def test_invalid_extension_attr(self):
    with self.assertRaisesRegex(ValueError,
                                'DmEnvAdaptor already has attribute'):
      dm_env_adaptor.DmEnvAdaptor(
          connection=mock.MagicMock(),
          specs=_SAMPLE_SPEC,
          extensions={'_connection': object()})


class CreateJoinHelpers(parameterized.TestCase):

  def test_create_world(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(
        return_value=dm_env_rpc_pb2.CreateWorldResponse(
            world_name='Damogran_01'))

    world_name = dm_env_adaptor.create_world(connection, {'planet': 'Damogran'})
    self.assertEqual('Damogran_01', world_name)

    connection.send.assert_called_once_with(
        text_format.Parse(
            """settings: {
                key: 'planet', value: { strings: { array: 'Damogran' } }
            }""", dm_env_rpc_pb2.CreateWorldRequest()))

  def test_join_world(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(
        return_value=dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_SPEC))

    env = dm_env_adaptor.join_world(connection, 'Damogran_01',
                                    {'player': 'zaphod'})
    self.assertIsNotNone(env)

    connection.send.assert_called_once_with(
        text_format.Parse(
            """world_name: 'Damogran_01'
                settings: {
                    key: 'player', value: { strings: { array: 'zaphod' } }
                }""", dm_env_rpc_pb2.JoinWorldRequest()))

  def test_create_join_world(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(side_effect=[
        dm_env_rpc_pb2.CreateWorldResponse(world_name='Damogran_01'),
        dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_SPEC)
    ])
    env, world_name = dm_env_adaptor.create_and_join_world(
        connection,
        create_world_settings={'planet': 'Damogran'},
        join_world_settings={
            'ship_type': 1,
            'player': 'zaphod',
        },
        requested_observations=['foo'])
    self.assertIsNotNone(env)
    self.assertSameElements(env.observation_spec().keys(), ['foo'])
    self.assertEqual('Damogran_01', world_name)

    connection.send.assert_has_calls([
        mock.call(
            text_format.Parse(
                """settings: {
                key: 'planet', value: { strings: { array: 'Damogran' } }
            }""", dm_env_rpc_pb2.CreateWorldRequest())),
        mock.call(
            text_format.Parse(
                """world_name: 'Damogran_01'
                settings: { key: 'ship_type', value: { int64s: { array: 1 } } }
                settings: {
                    key: 'player', value: { strings: { array: 'zaphod' } }
                }""", dm_env_rpc_pb2.JoinWorldRequest())),
    ])

  @parameterized.named_parameters(
      (
          'nested',
          {'nested': {'planet': 'Damogran'}},
          {'nested': {'ship_type': 1, 'player': 'zaphod'}},
      ),
      (
          'already_flattened',
          {'nested.planet': 'Damogran'},
          {'nested.ship_type': 1, 'nested.player': 'zaphod'},
      ),
  )
  def test_flatten_create_join_world_settings(
      self, create_settings, join_settings
  ):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(side_effect=[
        dm_env_rpc_pb2.CreateWorldResponse(world_name='Damogran_01'),
        dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_SPEC)
    ])
    env, world_name = dm_env_adaptor.create_and_join_world(
        connection,
        create_world_settings=create_settings,
        join_world_settings=join_settings)
    self.assertIsNotNone(env)
    self.assertEqual('Damogran_01', world_name)

    connection.send.assert_has_calls([
        mock.call(
            text_format.Parse(
                """settings: {
                key: 'nested.planet', value: { strings: { array: 'Damogran' } }
            }""", dm_env_rpc_pb2.CreateWorldRequest())),
        mock.call(
            text_format.Parse(
                """world_name: 'Damogran_01'
                settings: {
                  key: 'nested.ship_type', value: { int64s: { array: 1 } }
                }
                settings: {
                    key: 'nested.player', value: {
                        strings: { array: 'zaphod' }
                    }
                }""", dm_env_rpc_pb2.JoinWorldRequest())),
    ])

  def test_create_join_world_with_packed_settings(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(side_effect=[
        dm_env_rpc_pb2.CreateWorldResponse(world_name='Magrathea_02'),
        dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_SPEC)
    ])
    env_and_world_name = dm_env_adaptor.create_and_join_world(
        connection,
        create_world_settings={'planet': tensor_utils.pack_tensor('Magrathea')},
        join_world_settings={
            'ship_type': tensor_utils.pack_tensor(2),
            'player': tensor_utils.pack_tensor('arthur'),
            'unpacked_setting': [1, 2, 3],
        })
    self.assertIsNotNone(env_and_world_name.env)
    self.assertEqual('Magrathea_02', env_and_world_name.world_name)

    connection.send.assert_has_calls([
        mock.call(
            text_format.Parse(
                """settings: {
                key: 'planet', value: { strings: { array: 'Magrathea' } }
            }""", dm_env_rpc_pb2.CreateWorldRequest())),
        mock.call(
            text_format.Parse(
                """world_name: 'Magrathea_02'
                settings: { key: 'ship_type', value: { int64s: { array: 2 } } }
                settings: {
                    key: 'player', value: { strings: { array: 'arthur' } }
                }
                settings: {
                    key: 'unpacked_setting', value: {
                      int64s: { array: 1 array: 2 array: 3 }
                      shape: 3
                    }
                }""", dm_env_rpc_pb2.JoinWorldRequest())),
    ])

  def test_create_join_world_with_extension(self):

    class _ExampleExtension:

      def foo(self):
        return 'bar'

    connection = mock.MagicMock()
    connection.send = mock.MagicMock(side_effect=[
        dm_env_rpc_pb2.CreateWorldResponse(world_name='foo'),
        dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_SPEC)
    ])

    env, _ = dm_env_adaptor.create_and_join_world(
        connection,
        create_world_settings={},
        join_world_settings={},
        extensions={'extension': _ExampleExtension()})
    self.assertEqual('bar', env.extension.foo())

  def test_create_join_world_with_unnested_tensors(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(side_effect=[
        dm_env_rpc_pb2.CreateWorldResponse(world_name='Damogran_01'),
        dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_NESTED_SPECS)
    ])
    env, _ = dm_env_adaptor.create_and_join_world(
        connection,
        create_world_settings={},
        join_world_settings={},
        nested_tensors=False)
    expected_actions = {
        'foo.bar': specs.Array(shape=(), dtype=np.int32, name='foo.bar'),
        'baz': specs.Array(shape=(), dtype=str, name='baz'),
    }
    expected_observations = {
        'foo.bar': specs.Array(shape=(), dtype=np.int32, name='foo.bar'),
        'baz': specs.Array(shape=(), dtype=str, name='baz'),
    }

    self.assertSameElements(expected_actions, env.action_spec())
    self.assertSameElements(expected_observations, env.observation_spec())

  def test_create_join_world_with_invalid_extension(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(side_effect=[
        dm_env_rpc_pb2.CreateWorldResponse(world_name='foo'),
        dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_SPEC),
        dm_env_rpc_pb2.LeaveWorldResponse(),
        dm_env_rpc_pb2.DestroyWorldRequest()
    ])

    with self.assertRaisesRegex(ValueError,
                                'DmEnvAdaptor already has attribute'):
      _ = dm_env_adaptor.create_and_join_world(
          connection,
          create_world_settings={},
          join_world_settings={},
          extensions={'step': object()})

    connection.send.assert_has_calls([
        mock.call(dm_env_rpc_pb2.CreateWorldRequest()),
        mock.call(dm_env_rpc_pb2.JoinWorldRequest(world_name='foo')),
        mock.call(dm_env_rpc_pb2.LeaveWorldRequest()),
        mock.call(dm_env_rpc_pb2.DestroyWorldRequest(world_name='foo'))
    ])

  def test_created_but_failed_to_join_world(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(
        side_effect=(
            dm_env_rpc_pb2.CreateWorldResponse(world_name='Damogran_01'),
            error.DmEnvRpcError(status_pb2.Status(message='Failed to Join.')),
            dm_env_rpc_pb2.DestroyWorldResponse()))

    with self.assertRaisesRegex(error.DmEnvRpcError, 'Failed to Join'):
      _ = dm_env_adaptor.create_and_join_world(
          connection, create_world_settings={}, join_world_settings={})

    connection.send.assert_has_calls([
        mock.call(dm_env_rpc_pb2.CreateWorldRequest()),
        mock.call(dm_env_rpc_pb2.JoinWorldRequest(world_name='Damogran_01')),
        mock.call(dm_env_rpc_pb2.DestroyWorldRequest(world_name='Damogran_01'))
    ])

  def test_created_and_joined_but_adaptor_failed(self):
    connection = mock.MagicMock()
    connection.send = mock.MagicMock(
        side_effect=(
            dm_env_rpc_pb2.CreateWorldResponse(world_name='Damogran_01'),
            dm_env_rpc_pb2.JoinWorldResponse(specs=_SAMPLE_SPEC),
            dm_env_rpc_pb2.LeaveWorldResponse(),
            dm_env_rpc_pb2.DestroyWorldResponse()))

    with self.assertRaisesRegex(ValueError, 'Unsupported observations'):
      _ = dm_env_adaptor.create_and_join_world(
          connection,
          create_world_settings={},
          join_world_settings={},
          requested_observations=['invalid_observation'])

    connection.send.assert_has_calls([
        mock.call(dm_env_rpc_pb2.CreateWorldRequest()),
        mock.call(dm_env_rpc_pb2.JoinWorldRequest(world_name='Damogran_01')),
        mock.call(dm_env_rpc_pb2.LeaveWorldRequest()),
        mock.call(dm_env_rpc_pb2.DestroyWorldRequest(world_name='Damogran_01'))
    ])


if __name__ == '__main__':
  absltest.main()
