fix: flaky test

projectmesa · Aug 24, 2024 · cf1f0bb · cf1f0bb
1 parent 89b793b
commit cf1f0bb
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/polars/test_grid_polars.py b/tests/polars/test_grid_polars.py
@@ -156,9 +156,9 @@ def test_get_directions(
         grid_moore.place_agents(fix2_AgentSetPolars, [[0, 1], [0, 2], [1, 0], [1, 2]])
         agents = grid_moore.agents.sort("agent_id")
 
-        assert grid_moore.agents["agent_id"].to_list() == [0, 1, 4, 5, 6, 7]
-        assert grid_moore.agents["dim_0"].to_list() == [0, 1, 0, 0, 1, 1]
-        assert grid_moore.agents["dim_1"].to_list() == [0, 1, 1, 2, 0, 2]
+        assert agents["agent_id"].to_list() == [0, 1, 4, 5, 6, 7]
+        assert agents["dim_0"].to_list() == [0, 1, 0, 0, 1, 1]
+        assert agents["dim_1"].to_list() == [0, 1, 1, 2, 0, 2]
         dir = grid_moore.get_directions(agents0=[0, 1], agents1=[4, 5])
         print(dir)
         assert isinstance(dir, pl.DataFrame)