test_datacollector.py

"""
Test the DataCollector
"""
import unittest

from mesa import Model, Agent
from mesa.time import BaseScheduler
from mesa.datacollection import DataCollector


class MockAgent(Agent):
    """
    Minimalistic agent for testing purposes.
    """

    def __init__(self, unique_id, model, val=0):
        super().__init__(unique_id, model)
        self.val = val
        self.val2 = val

    def step(self):
        """
        Increment vals by 1.
        """
        self.val += 1
        self.val2 += 1

    def write_final_values(self):
        """
        Write the final value to the appropriate table.
        """
        row = {"agent_id": self.unique_id, "final_value": self.val}
        self.model.datacollector.add_table_row("Final_Values", row)


class MockModel(Model):
    """
    Minimalistic model for testing purposes.
    """

    schedule = BaseScheduler(None)

    def __init__(self):
        self.schedule = BaseScheduler(self)
        self.model_val = 100

        for i in range(10):
            a = MockAgent(i, self, val=i)
            self.schedule.add(a)
        self.datacollector = DataCollector(
            {
                "total_agents": lambda m: m.schedule.get_agent_count(),
                "model_value": "model_val",
                "model_calc": self.schedule.get_agent_count,
                "model_calc_comp": [self.test_model_calc_comp, [3, 4]],
                "model_calc_fail": [self.test_model_calc_comp, [12, 0]],
            },
            {"value": lambda a: a.val, "value2": "val2"},
            {"Final_Values": ["agent_id", "final_value"]},
        )

    def test_model_calc_comp(self, input1, input2):
        if input2 > 0:
            return (self.model_val * input1) / input2
        else:
            assert ValueError
            return None

    def step(self):
        self.schedule.step()
        self.datacollector.collect(self)


class TestDataCollector(unittest.TestCase):
    def setUp(self):
        """
        Create the model and run it a set number of steps.
        """
        self.model = MockModel()
        for i in range(7):
            if i == 4:
                self.model.schedule.remove(self.model.schedule._agents[3])
            self.model.step()

        # Write to table:
        for agent in self.model.schedule.agents:
            agent.write_final_values()

    def step_assertion(self, model_var):
        for element in model_var:
            if model_var.index(element) < 4:
                assert element == 10
            else:
                assert element == 9

    def test_model_vars(self):
        """
        Test model-level variable collection.
        """
        data_collector = self.model.datacollector
        assert "total_agents" in data_collector.model_vars
        assert "model_value" in data_collector.model_vars
        assert "model_calc" in data_collector.model_vars
        assert "model_calc_comp" in data_collector.model_vars
        assert "model_calc_fail" in data_collector.model_vars
        assert len(data_collector.model_vars["total_agents"]) == 7
        assert len(data_collector.model_vars["model_value"]) == 7
        assert len(data_collector.model_vars["model_calc"]) == 7
        assert len(data_collector.model_vars["model_calc_comp"]) == 7
        self.step_assertion(data_collector.model_vars["total_agents"])
        for element in data_collector.model_vars["model_value"]:
            assert element == 100
        self.step_assertion(data_collector.model_vars["model_calc"])
        for element in data_collector.model_vars["model_calc_comp"]:
            assert element == 75
        for element in data_collector.model_vars["model_calc_fail"]:
            assert element is None

    def test_agent_records(self):
        """
        Test agent-level variable collection.
        """
        data_collector = self.model.datacollector
        agent_table = data_collector.get_agent_vars_dataframe()

        assert len(data_collector._agent_records) == 7
        for step, records in data_collector._agent_records.items():
            if step < 5:
                assert len(records) == 10
            else:
                assert len(records) == 9

            for values in records:
                assert len(values) == 4

        assert "value" in list(agent_table.columns)
        assert "value2" in list(agent_table.columns)
        assert "value3" not in list(agent_table.columns)

        with self.assertRaises(KeyError):
            data_collector._agent_records[8]

    def test_table_rows(self):
        """
        Test table collection
        """
        data_collector = self.model.datacollector
        assert len(data_collector.tables["Final_Values"]) == 2
        assert "agent_id" in data_collector.tables["Final_Values"]
        assert "final_value" in data_collector.tables["Final_Values"]
        for key, data in data_collector.tables["Final_Values"].items():
            assert len(data) == 9

        with self.assertRaises(Exception):
            data_collector.add_table_row("error_table", {})

        with self.assertRaises(Exception):
            data_collector.add_table_row("Final_Values", {"final_value": 10})

    def test_exports(self):
        """
        Test DataFrame exports
        """
        data_collector = self.model.datacollector
        model_vars = data_collector.get_model_vars_dataframe()
        agent_vars = data_collector.get_agent_vars_dataframe()
        table_df = data_collector.get_table_dataframe("Final_Values")
        assert model_vars.shape == (7, 5)
        assert agent_vars.shape == (67, 2)
        assert table_df.shape == (9, 2)

        with self.assertRaises(Exception):
            table_df = data_collector.get_table_dataframe("not a real table")


if __name__ == "__main__":
    unittest.main()