# Path: meltingpot/python/utils/substrates/wrappers/multiplayer_wrapper.py
# T = TypeVar("T")
# def _player_observations(observations: Mapping[str, T], suffix: str,
#                          num_players: int) -> Iterator[T]:
#   def __init__(self, env,
#                individual_observation_names: Sequence[str],
#                global_observation_names: Sequence[str]):
#   def _get_num_players(self) -> int:
#   def _get_observations(
#       self, source: Mapping[str, T]) -> Sequence[Mapping[str, T]]:
#   def _get_rewards(self, source: Mapping[str, T]) -> Sequence[T]:
#   def _get_timestep(self, source: dm_env.TimeStep) -> dm_env.TimeStep:
#   def _get_action(self, source: Sequence[Mapping[str, T]]) -> Mapping[str, T]:
#   def reset(self) -> dm_env.TimeStep:
#   def step(
#       self, actions: Sequence[Mapping[str, np.ndarray]]) -> dm_env.TimeStep:
#   def action_spec(self) -> Sequence[Mapping[str, dm_env.specs.Array]]:
#   def observation_spec(self) -> Sequence[Mapping[str, dm_env.specs.Array]]:
#   def reward_spec(self) -> Sequence[dm_env.specs.Array]:
# class Wrapper(base.Wrapper):


from unittest import mock
from absl.testing import absltest
from meltingpot.python.utils.substrates.wrappers import multiplayer_wrapper
import dm_env
import numpy as np
import dmlab2d
#
# 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 multiplayer_wrapper."""





ACT_SPEC = dm_env.specs.BoundedArray(
    shape=(), minimum=0, maximum=4, dtype=np.int8)
ACT_VALUE = np.ones([], dtype=np.int8)
RGB_SPEC = dm_env.specs.Array(shape=(8, 8, 3), dtype=np.int8)
RGB_VALUE = np.ones((8, 8, 3), np.int8)
REWARD_SPEC = dm_env.specs.Array(shape=(), dtype=np.float32)
REWARD_VALUE = np.ones((), dtype=np.float32)


class Lab2DToListsWrapperTest(absltest.TestCase):

  def test_get_num_players(self):
    env = mock.Mock(spec_set=dmlab2d.Environment)
    env.action_spec.return_value = {
        '1.MOVE': ACT_SPEC,
        '2.MOVE': ACT_SPEC,
        '3.MOVE': ACT_SPEC
    }

# Path: meltingpot/python/utils/substrates/wrappers/multiplayer_wrapper.py
# T = TypeVar("T")
# def _player_observations(observations: Mapping[str, T], suffix: str,
#                          num_players: int) -> Iterator[T]:
#   def __init__(self, env,
#                individual_observation_names: Sequence[str],
#                global_observation_names: Sequence[str]):
#   def _get_num_players(self) -> int:
#   def _get_observations(
#       self, source: Mapping[str, T]) -> Sequence[Mapping[str, T]]:
#   def _get_rewards(self, source: Mapping[str, T]) -> Sequence[T]:
#   def _get_timestep(self, source: dm_env.TimeStep) -> dm_env.TimeStep:
#   def _get_action(self, source: Sequence[Mapping[str, T]]) -> Mapping[str, T]:
#   def reset(self) -> dm_env.TimeStep:
#   def step(
#       self, actions: Sequence[Mapping[str, np.ndarray]]) -> dm_env.TimeStep:
#   def action_spec(self) -> Sequence[Mapping[str, dm_env.specs.Array]]:
#   def observation_spec(self) -> Sequence[Mapping[str, dm_env.specs.Array]]:
#   def reward_spec(self) -> Sequence[dm_env.specs.Array]:
# class Wrapper(base.Wrapper):


# Path: meltingpot/python/utils/substrates/wrappers/multiplayer_wrapper_test.py
from unittest import mock
from absl.testing import absltest
from meltingpot.python.utils.substrates.wrappers import multiplayer_wrapper
import dm_env
import numpy as np
import dmlab2d

#
# 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 multiplayer_wrapper."""





ACT_SPEC = dm_env.specs.BoundedArray(
    shape=(), minimum=0, maximum=4, dtype=np.int8)
ACT_VALUE = np.ones([], dtype=np.int8)
RGB_SPEC = dm_env.specs.Array(shape=(8, 8, 3), dtype=np.int8)
RGB_VALUE = np.ones((8, 8, 3), np.int8)
REWARD_SPEC = dm_env.specs.Array(shape=(), dtype=np.float32)
REWARD_VALUE = np.ones((), dtype=np.float32)


class Lab2DToListsWrapperTest(absltest.TestCase):

  def test_get_num_players(self):
    env = mock.Mock(spec_set=dmlab2d.Environment)
    env.action_spec.return_value = {
        '1.MOVE': ACT_SPEC,
        '2.MOVE': ACT_SPEC,
        '3.MOVE': ACT_SPEC
    }

    wrapped = multiplayer_wrapper.Wrapper(