fix: update stale tests (provider key + streaming mock) and ruff format

Author: DeveshParagiri

extropy/population/sampler/core.py

@@ -171,7 +171,6 @@ class SamplingError(Exception):
     pass
 
 
-
 def _has_household_attributes(spec: PopulationSpec) -> bool:
     """Check if the spec has household-scoped attributes, indicating household mode."""
     return any(attr.scope == "household" for attr in spec.attributes)
@@ -452,7 +451,11 @@ def _sample_population_households(
     """
     if config is None:
         config = HouseholdConfig()
-    focus_mode = agent_focus_mode if agent_focus_mode in ("all", "couples", "primary_only") else "primary_only"
+    focus_mode = (
+        agent_focus_mode
+        if agent_focus_mode in ("all", "couples", "primary_only")
+        else "primary_only"
+    )
 
     hh_id_width = len(str(target_n - 1))  # safe upper bound for household IDs
 

tests/test_agent_focus.py

@@ -140,7 +140,9 @@ class TestAgentFocusPrimaryOnly:
 
     def test_partners_are_npc(self):
         spec = _make_household_spec()
-        result = sample_population(spec, count=200, seed=42, agent_focus_mode="primary_only")
+        result = sample_population(
+            spec, count=200, seed=42, agent_focus_mode="primary_only"
+        )
         agents = result.agents
 
         # Find primary adults with partners
@@ -168,7 +170,9 @@ def test_partners_are_npc(self):
 
     def test_no_partner_agents_in_result(self):
         spec = _make_household_spec()
-        result = sample_population(spec, count=200, seed=42, agent_focus_mode="primary_only")
+        result = sample_population(
+            spec, count=200, seed=42, agent_focus_mode="primary_only"
+        )
         agents = result.agents
 
         # Count agents with household_role = adult_secondary
@@ -183,12 +187,16 @@ def test_no_partner_agents_in_result(self):
     def test_exact_agent_count(self):
         """Regression test: -n must produce exactly N agents."""
         spec = _make_household_spec()
-        result = sample_population(spec, count=500, seed=42, agent_focus_mode="primary_only")
+        result = sample_population(
+            spec, count=500, seed=42, agent_focus_mode="primary_only"
+        )
         assert len(result.agents) == 500
 
     def test_npc_partner_has_correlated_demographics(self):
         spec = _make_household_spec()
-        result = sample_population(spec, count=300, seed=42, agent_focus_mode="primary_only")
+        result = sample_population(
+            spec, count=300, seed=42, agent_focus_mode="primary_only"
+        )
         agents = result.agents
 
         age_diffs = []
@@ -211,7 +219,9 @@ def test_npc_partner_has_correlated_demographics(self):
 
     def test_npc_partner_shares_last_name(self):
         spec = _make_household_spec()
-        result = sample_population(spec, count=200, seed=42, agent_focus_mode="primary_only")
+        result = sample_population(
+            spec, count=200, seed=42, agent_focus_mode="primary_only"
+        )
 
         for agent in result.agents:
             npc = agent.get("partner_npc")

tests/test_estimator.py

@@ -217,10 +217,8 @@ def test_resolve_openai_default(self):
         assert resolve_default_model("openai", "simple") == "gpt-5-mini"
 
     def test_resolve_claude_default(self):
-        assert (
-            resolve_default_model("claude", "reasoning") == "claude-sonnet-4-5-20250929"
-        )
-        assert resolve_default_model("claude", "simple") == "claude-haiku-4-5-20251001"
+        assert resolve_default_model("anthropic", "reasoning") == "claude-sonnet-4-6"
+        assert resolve_default_model("anthropic", "simple") == "claude-haiku-4-5-20251001"
 
     def test_resolve_unknown_provider_falls_back(self):
         model = resolve_default_model("unknown_provider", "reasoning")

tests/test_providers.py

@@ -404,7 +404,12 @@ def test_extracts_sources(self, mock_get_client):
         response.content = [search_block, tool_block, text_block]
 
         mock_client = MagicMock()
-        mock_client.messages.create.return_value = response
+        # Code uses client.messages.stream() context manager → .get_final_message()
+        mock_stream = MagicMock()
+        mock_stream.__enter__ = MagicMock(return_value=mock_stream)
+        mock_stream.__exit__ = MagicMock(return_value=False)
+        mock_stream.get_final_message.return_value = response
+        mock_client.messages.stream.return_value = mock_stream
         mock_get_client.return_value = mock_client
 
         result, sources = provider.agentic_research(