Fix assertions for index

Author: luca-patrignani

tests/pandas/test_agentset_pandas.py

@@ -11,7 +11,7 @@
 
 @tg.typechecked
 class ExampleAgentSetPandas(AgentSetPandas):
-    def __init__(self, model: ModelDF):
+    def __init__(self, model: ModelDF, index: pd.Index):
         super().__init__(model)
         self.starting_wealth = pd.Series([1, 2, 3, 4], name="wealth")
 
@@ -25,7 +25,7 @@ def step(self) -> None:
 @pytest.fixture
 def fix1_AgentSetPandas() -> ExampleAgentSetPandas:
     model = ModelDF()
-    agents = ExampleAgentSetPandas(model)
+    agents = ExampleAgentSetPandas(model, pd.Index([0, 1, 2, 3], name="unique_id"))
     agents.add({"unique_id": [0, 1, 2, 3]})
     agents["wealth"] = agents.starting_wealth
     agents["age"] = [10, 20, 30, 40]
@@ -36,7 +36,7 @@ def fix1_AgentSetPandas() -> ExampleAgentSetPandas:
 @pytest.fixture
 def fix2_AgentSetPandas() -> ExampleAgentSetPandas:
     model = ModelDF()
-    agents = ExampleAgentSetPandas(model)
+    agents = ExampleAgentSetPandas(model, pd.Index([4, 5, 6, 7], name="unique_id"))
     agents.add({"unique_id": [4, 5, 6, 7]})
     agents["wealth"] = agents.starting_wealth + 10
     agents["age"] = [100, 200, 300, 400]
@@ -55,7 +55,7 @@ def fix1_AgentSetPandas_with_pos(fix1_AgentSetPandas) -> ExampleAgentSetPandas:
 class Test_AgentSetPandas:
     def test__init__(self):
         model = ModelDF()
-        agents = ExampleAgentSetPandas(model)
+        agents = ExampleAgentSetPandas(model, pd.Index([0, 1, 2, 3]))
         assert agents.model == model
         assert isinstance(agents.agents, pd.DataFrame)
         assert agents.agents.index.name == "unique_id"
@@ -78,14 +78,15 @@ def test_add(
 
         # Test with a list (Sequence[Any])
         result = agents.add([10, 5, 10], inplace=False)
-        assert result.agents.index.to_list() == [0, 1, 2, 3, 10]
+        assert result.agents.index.to_list() == [0, 1, 2, 3, 4]
         assert result.agents.wealth.to_list() == [1, 2, 3, 4, 5]
         assert result.agents.age.to_list() == [10, 20, 30, 40, 10]
         assert agents.agents.index.name == "unique_id"
 
         # Test with a dict[str, Any]
         agents.add({"unique_id": [4, 5], "wealth": [5, 6], "age": [50, 60]})
         assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5, 6]
+        assert agents.agents.index.tolist() == [0, 1, 2, 3, 4, 5]
         assert agents.agents.age.tolist() == [10, 20, 30, 40, 50, 60]
         assert agents.agents.index.name == "unique_id"
 
@@ -286,6 +287,7 @@ def test__add__(
 
         # Test with an AgentSetPandas and a dict
         agents3 = agents + {"unique_id": 10, "wealth": 5}
+        assert agents3.agents.index.tolist() == [0, 1, 2, 3, 4]
         assert agents3.agents.wealth.tolist() == [1, 2, 3, 4, 5]
 
     def test__contains__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
@@ -358,6 +360,7 @@ def test__iadd__(
         # Test with an AgentSetPandas and a dict
         agents = deepcopy(fix1_AgentSetPandas)
         agents += {"unique_id": 10, "wealth": 5}
+        assert agents.agents.index.tolist() == [0, 1, 2, 3, 4]
         assert agents.agents.wealth.tolist() == [1, 2, 3, 4, 5]
 
     def test__iter__(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
@@ -436,24 +439,25 @@ def test_agents(
 
         # Test agents.setter
         agents.agents = agents2.agents
-        assert len(agents.active_agents) == 4
+        assert agents.agents.wealth.tolist() == [11, 12, 13, 14]
+        assert agents.agents.age.tolist() == [100, 200, 300, 400]
 
     def test_active_agents(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         # Test with select
         agents.select(agents["wealth"] > 2, inplace=True)
-        assert len(agents.active_agents) == 2
+        assert agents.active_agents.index.tolist() == [2, 3]
 
         # Test with active_agents.setter
         agents.active_agents = agents.agents.wealth > 2
-        assert len(agents.active_agents) == 2
+        assert agents.active_agents.index.to_list() == [2, 3]
 
     def test_inactive_agents(self, fix1_AgentSetPandas: ExampleAgentSetPandas):
         agents = fix1_AgentSetPandas
 
         agents.select(agents["wealth"] > 2, inplace=True)
-        assert len(agents.active_agents) == 2
+        assert agents.inactive_agents.index.to_list() == [0, 1]
 
     def test_pos(self, fix1_AgentSetPandas_with_pos: ExampleAgentSetPandas):
         pos = fix1_AgentSetPandas_with_pos.pos