fix test_examples failure

pahulgogna · Aug 30, 2020 · 230590d · 230590d
1 parent 899ce22
commit 230590d
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 3 deletions.
diff --git a/mesa/datacollection.py b/mesa/datacollection.py
@@ -247,6 +247,7 @@ def get_agent_vars_dataframe(self):
         df = pd.DataFrame.from_records(
             data=all_records, columns=["Step", "AgentID"] + rep_names,
         )
+        df = df.set_index(["Step", "AgentID"])
         return df
 
     def get_table_dataframe(self, table_name):

diff --git a/tests/test_batchrunner.py b/tests/test_batchrunner.py
@@ -144,7 +144,7 @@ def test_agent_level_vars(self):
         assert "agent_val" in list(agent_vars.columns)
         assert "val_non_existent" not in list(agent_vars.columns)
         assert "agent_id" in list(agent_collector[(0, 1, 1)].columns)
-        assert "Step" in list(agent_collector[(0, 1, 5)].columns)
+        assert "Step" in list(agent_collector[(0, 1, 5)].index)
         assert "nose" not in list(agent_collector[(0, 1, 1)].columns)
 
         self.assertEqual(

diff --git a/tests/test_batchrunnerMP.py b/tests/test_batchrunnerMP.py
@@ -147,7 +147,7 @@ def test_agent_level_vars(self):
         assert "agent_val" in list(agent_vars.columns)
         assert "val_non_existent" not in list(agent_vars.columns)
         assert "agent_id" in list(agent_collector[(0, 1, 1)].columns)
-        assert "Step" in list(agent_collector[(0, 1, 5)].columns)
+        assert "Step" in list(agent_collector[(0, 1, 5)].index)
         assert "nose" not in list(agent_collector[(0, 1, 1)].columns)
 
         self.assertEqual(

diff --git a/tests/test_datacollector.py b/tests/test_datacollector.py
@@ -154,7 +154,7 @@ def test_exports(self):
         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 == (70, 4)
+        assert agent_vars.shape == (70, 2)
         assert table_df.shape == (10, 2)
 
         with self.assertRaises(Exception):