/
test_batchrunner.py
94 lines (82 loc) · 2.51 KB
/
test_batchrunner.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
"""
Test the BatchRunner
"""
import unittest
from mesa import Agent, Model
from mesa.batchrunner import BatchRunner
from mesa.time import BaseScheduler
NUM_AGENTS = 7
class MockAgent(Agent):
"""
Minimalistic model for testing purposes
"""
def __init__(self, unique_id, val):
self.unique_id = unique_id
self.val = val
def step(self):
"""
increment val by 1
"""
self.val += 1
class MockModel(Model):
"""
Minimalistic model for testing purposes
"""
def __init__(self, model_param, agent_param):
"""
Args:
model_param (any): parameter specific to the model
agent_param (int): parameter specific to the agent
"""
self.schedule = BaseScheduler(None)
self.model_param = model_param
self.running = True
for i in range(NUM_AGENTS):
a = MockAgent(i, agent_param)
self.schedule.add(a)
def step(self):
self.schedule.step()
class TestBatchRunner(unittest.TestCase):
"""
Test that BatchRunner is running batches
"""
def setUp(self):
"""
Create the model and run it for some steps
"""
self.model_reporter = {"model": lambda m: m.model_param}
self.agent_reporter = {
"agent_id": lambda a: a.unique_id,
"agent_val": lambda a: a.val}
self.params = {
'model_param': range(3),
'agent_param': [1, 8],
}
self.iterations = 17
self.batch = BatchRunner(
MockModel,
self.params,
iterations=self.iterations,
max_steps=3,
model_reporters=self.model_reporter,
agent_reporters=self.agent_reporter)
self.batch.run_all()
def test_model_level_vars(self):
"""
Test that model-level variable collection is of the correct size
"""
model_vars = self.batch.get_model_vars_dataframe()
rows = len(self.params['model_param']) * \
len(self.params['agent_param']) * \
self.iterations
assert model_vars.shape == (rows, 4)
def test_agent_level_vars(self):
"""
Test that agent-level variable collection is of the correct size
"""
agent_vars = self.batch.get_agent_vars_dataframe()
rows = NUM_AGENTS * \
len(self.params['agent_param']) * \
len(self.params['model_param']) * \
self.iterations
assert agent_vars.shape == (rows, 6)