from pymasep.common import Controller, State, Agent
from test.unit.mocks.game_mock import GameMock


class ControllerMock(Controller):

    def action_choice(self, observation: State, agent: Agent):
        """
        Choose the action. (Default, at the moment :  Action.ACTION_TYPE_2)
        :param observation the observation used to choose the action
        :param agent the agent who choose the action
        :return the action chosen for the agent.
        """
        return self.environment.create_action(GameMock.ACTION_TYPE_2)
