# Path: meltingpot/python/utils/scenarios/wrappers/base.py
# class Wrapper(substrate.Substrate):
#   def __init__(self, env: substrate.Substrate):
#   def reset(self, *args, **kwargs) -> ...:
#   def step(self, *args, **kwargs) -> ...:
#   def reward_spec(self, *args, **kwargs) -> ...:
#   def discount_spec(self, *args, **kwargs) -> ...:
#   def observation_spec(self, *args, **kwargs) -> ...:
#   def action_spec(self, *args, **kwargs) -> ...:
#   def close(self, *args, **kwargs) -> ...:
#   def observation(self, *args, **kwargs) -> ...:
#   def events(self, *args, **kwargs) -> ...:
#   def list_property(self, *args, **kwargs) -> ...:
#   def write_property(self, *args, **kwargs) -> ...:
#   def read_property(self, *args, **kwargs) -> ...:
#   def observables(self, *args, **kwargs) -> ...:
# 
# Path: meltingpot/python/utils/scenarios/wrappers/default_observation_wrapper.py
# def _setdefault(dictionary: Dict[str, Any], key: str,
#                 value: Any) -> Dict[str, Any]:
#   def __init__(self,
#                env: base.Substrate,
#                key: str,
#                default_value: np.ndarray,
#                default_spec: Optional[dm_env.specs.Array] = None):
#   def reset(self):
#   def step(self, action):
#   def observation_spec(self):
# class Wrapper(base.Wrapper):


from unittest import mock
from absl.testing import absltest
from meltingpot.python.utils.scenarios.wrappers import base
from meltingpot.python.utils.scenarios.wrappers import default_observation_wrapper
import dm_env
import numpy as np
"""Tests for default_observation_wrapper."""




TEST_KEY = 'test_key'
RGB_SPEC = dm_env.specs.Array(shape=(8, 8, 3), dtype=np.int8)
RGB_VALUE = np.ones((8, 8, 3), np.int8)
DEFAULT_SPEC = dm_env.specs.Array(shape=[0], dtype=np.double, name=TEST_KEY)
DEFAULT_VALUE = np.zeros(shape=[0], dtype=np.double)
TEST_SPEC = dm_env.specs.Array(shape=[2], dtype=np.double, name=TEST_KEY)
TEST_VALUE = np.zeros(shape=[2], dtype=np.double)


class WrapperTest(absltest.TestCase):

  def test_setdefault_no_change(self):
    original = {'a': 0}
    actual = default_observation_wrapper._setdefault(original, 'a', 100)
    self.assertIs(actual, original)

  def test_setdefault_change(self):
    original = {'a': 0}
    actual = default_observation_wrapper._setdefault(original, 'b', 100)
    with self.subTest('creates_copy'):
      self.assertIsNot(actual, original)
    with self.subTest('adds_value'):
      self.assertEqual(actual, {'a': 0, 'b': 100})

  def test_change(self):

# Path: meltingpot/python/utils/scenarios/wrappers/base.py
# class Wrapper(substrate.Substrate):
#   def __init__(self, env: substrate.Substrate):
#   def reset(self, *args, **kwargs) -> ...:
#   def step(self, *args, **kwargs) -> ...:
#   def reward_spec(self, *args, **kwargs) -> ...:
#   def discount_spec(self, *args, **kwargs) -> ...:
#   def observation_spec(self, *args, **kwargs) -> ...:
#   def action_spec(self, *args, **kwargs) -> ...:
#   def close(self, *args, **kwargs) -> ...:
#   def observation(self, *args, **kwargs) -> ...:
#   def events(self, *args, **kwargs) -> ...:
#   def list_property(self, *args, **kwargs) -> ...:
#   def write_property(self, *args, **kwargs) -> ...:
#   def read_property(self, *args, **kwargs) -> ...:
#   def observables(self, *args, **kwargs) -> ...:
# 
# Path: meltingpot/python/utils/scenarios/wrappers/default_observation_wrapper.py
# def _setdefault(dictionary: Dict[str, Any], key: str,
#                 value: Any) -> Dict[str, Any]:
#   def __init__(self,
#                env: base.Substrate,
#                key: str,
#                default_value: np.ndarray,
#                default_spec: Optional[dm_env.specs.Array] = None):
#   def reset(self):
#   def step(self, action):
#   def observation_spec(self):
# class Wrapper(base.Wrapper):


# Path: meltingpot/python/utils/scenarios/wrappers/default_observation_wrapper_test.py
from unittest import mock
from absl.testing import absltest
from meltingpot.python.utils.scenarios.wrappers import base
from meltingpot.python.utils.scenarios.wrappers import default_observation_wrapper
import dm_env
import numpy as np

"""Tests for default_observation_wrapper."""




TEST_KEY = 'test_key'
RGB_SPEC = dm_env.specs.Array(shape=(8, 8, 3), dtype=np.int8)
RGB_VALUE = np.ones((8, 8, 3), np.int8)
DEFAULT_SPEC = dm_env.specs.Array(shape=[0], dtype=np.double, name=TEST_KEY)
DEFAULT_VALUE = np.zeros(shape=[0], dtype=np.double)
TEST_SPEC = dm_env.specs.Array(shape=[2], dtype=np.double, name=TEST_KEY)
TEST_VALUE = np.zeros(shape=[2], dtype=np.double)


class WrapperTest(absltest.TestCase):

  def test_setdefault_no_change(self):
    original = {'a': 0}
    actual = default_observation_wrapper._setdefault(original, 'a', 100)
    self.assertIs(actual, original)

  def test_setdefault_change(self):
    original = {'a': 0}
    actual = default_observation_wrapper._setdefault(original, 'b', 100)
    with self.subTest('creates_copy'):
      self.assertIsNot(actual, original)
    with self.subTest('adds_value'):
      self.assertEqual(actual, {'a': 0, 'b': 100})

  def test_change(self):

    env = mock.Mock(spec_set=base.Substrate)