1 import pytest

2 from ..agent_model.agents import AtmosphereEqualizerAgent

3 from ..agent_model.util import get_default_currency_data

4

5 class MockAgent:

6 def __init__(self, amount, volume, o2, co2):

7 self.amount = amount

8 self.properties = {'volume': {'value': volume}}

9 self.storage = {'o2': o2, 'co2': co2}

10 self.capacity = {'o2': 100, 'co2': 100}

11

12 def view(self, *args, **kwargs):

13 return {k: v for k, v in self.storage.items()}

14

15 def increment(self, currency, amount):

16 self.storage[currency] += amount

17

18 class MockModel:

19 agents = {}

20 currencies = get_default_currency_data()

21

22@pytest.fixture

23 def mock_model():

24 model = MockModel()

25 model.agents = {

26 'test_agent_1': MockAgent(1, 20, 3, 3),

27 'test_agent_2': MockAgent(1, 10, 3, 3),

28 }

29 return model

30

31 class TestAgentAtmosphereEqualizer:

32 def test_agent_atmosphere_equalizer_step(self, mock_model):

33 conns = ['test_agent_1', 'test_agent_2']

34 kwargs = {'agent_id': 'atmosphere_equalizer',

35 'flows': {'in': {'atmosphere': {'connections': conns}},

36 'out': {'atmosphere': {'connections': conns}}}}

37 agent = AtmosphereEqualizerAgent(mock_model, **kwargs)

38 agent.step()

39 test_agent_1 = mock_model.agents['test_agent_1']

40 test_agent_2 = mock_model.agents['test_agent_2']

41 assert test_agent_1.storage['o2'] == 4

42 assert test_agent_1.storage['co2'] == 4

43 assert test_agent_2.storage['o2'] == 2

44 assert test_agent_2.storage['co2'] == 2

45