# Path: meltingpot/python/utils/scenarios/wrappers/all_observations_wrapper.py
# GLOBAL_KEY = 'global'
# OBSERVATIONS_KEY = 'observations'
# REWARDS_KEY = 'rewards'
# ACTIONS_KEY = 'actions'
# class Wrapper(base.Wrapper):
#   def __init__(self, env: base.Substrate,
#                observations_to_share: Collection[str] = (),
#                share_actions: bool = False,
#                share_rewards: bool = False) -> None:
#   def _shared_observation(
#       self,
#       observations: Sequence[Mapping[str, np.ndarray]],
#       rewards: Sequence[Union[float, np.ndarray]],
#       actions: Sequence[int]):
#   def _adjusted_timestep(self, timestep: dm_env.TimeStep,
#                          actions: Sequence[int]) -> dm_env.TimeStep:
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, actions: Sequence[int]) -> dm_env.TimeStep:
#   def _shared_observation_spec(
#       self,
#       observation_spec: Mapping[str, dm_env.specs.Array],
#       reward_spec: dm_env.specs.Array,
#       action_spec: dm_env.specs.DiscreteArray):
#   def observation_spec(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) -> ...:


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from meltingpot.python.utils.scenarios.wrappers import all_observations_wrapper
from meltingpot.python.utils.scenarios.wrappers import base
import dm_env
import numpy as np
# You may obtain a copy of the License at
#
#     https://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 all_observations_wrapper."""




GLOBAL_KEY = all_observations_wrapper.GLOBAL_KEY
REWARDS_KEY = all_observations_wrapper.REWARDS_KEY
OBSERVATIONS_KEY = all_observations_wrapper.OBSERVATIONS_KEY
ACTIONS_KEY = all_observations_wrapper.ACTIONS_KEY

ACTION_1 = 'action_1'
ACTION_2 = 'action_2'
OBSERVATION_1 = 'observation_1'
OBSERVATION_2 = 'observation_2'
ACTION_SPEC = dm_env.specs.DiscreteArray(num_values=5, dtype=np.int32)
REWARD_SPEC = dm_env.specs.Array(shape=[], dtype=np.float32)


class AllObservationsWrapperTest(parameterized.TestCase):

  def test_observation_spec(self):

# Path: meltingpot/python/utils/scenarios/wrappers/all_observations_wrapper.py
# GLOBAL_KEY = 'global'
# OBSERVATIONS_KEY = 'observations'
# REWARDS_KEY = 'rewards'
# ACTIONS_KEY = 'actions'
# class Wrapper(base.Wrapper):
#   def __init__(self, env: base.Substrate,
#                observations_to_share: Collection[str] = (),
#                share_actions: bool = False,
#                share_rewards: bool = False) -> None:
#   def _shared_observation(
#       self,
#       observations: Sequence[Mapping[str, np.ndarray]],
#       rewards: Sequence[Union[float, np.ndarray]],
#       actions: Sequence[int]):
#   def _adjusted_timestep(self, timestep: dm_env.TimeStep,
#                          actions: Sequence[int]) -> dm_env.TimeStep:
#   def reset(self) -> dm_env.TimeStep:
#   def step(self, actions: Sequence[int]) -> dm_env.TimeStep:
#   def _shared_observation_spec(
#       self,
#       observation_spec: Mapping[str, dm_env.specs.Array],
#       reward_spec: dm_env.specs.Array,
#       action_spec: dm_env.specs.DiscreteArray):
#   def observation_spec(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/all_observations_wrapper_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from meltingpot.python.utils.scenarios.wrappers import all_observations_wrapper
from meltingpot.python.utils.scenarios.wrappers import base
import dm_env
import numpy as np

# You may obtain a copy of the License at
#
#     https://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 all_observations_wrapper."""




GLOBAL_KEY = all_observations_wrapper.GLOBAL_KEY
REWARDS_KEY = all_observations_wrapper.REWARDS_KEY
OBSERVATIONS_KEY = all_observations_wrapper.OBSERVATIONS_KEY
ACTIONS_KEY = all_observations_wrapper.ACTIONS_KEY

ACTION_1 = 'action_1'
ACTION_2 = 'action_2'
OBSERVATION_1 = 'observation_1'
OBSERVATION_2 = 'observation_2'
ACTION_SPEC = dm_env.specs.DiscreteArray(num_values=5, dtype=np.int32)
REWARD_SPEC = dm_env.specs.Array(shape=[], dtype=np.float32)


class AllObservationsWrapperTest(parameterized.TestCase):

  def test_observation_spec(self):

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