fix mypy issues

Author: cbrady

src/schematic/datastream/datastream_client.py

@@ -53,7 +53,7 @@ def _coerce_nulls(data: dict, model_cls: type) -> dict:
 def _validate(model_cls: type, raw: Any) -> Any:
     """Validate raw data into a Pydantic model, coercing Go-style nulls first."""
     if isinstance(raw, dict):
-        return model_cls.model_validate(_coerce_nulls(raw, model_cls))
+        return model_cls.model_validate(_coerce_nulls(raw, model_cls))  # type: ignore[attr-defined]
     return raw
 
 
@@ -396,8 +396,8 @@ async def check_flag(
         else:
             tasks.append(_resolved(cached_user))
 
-        company, user = await asyncio.gather(*tasks)
-        return self._evaluate_flag(flag, company, user)
+        results: list = await asyncio.gather(*tasks)
+        return self._evaluate_flag(flag, results[0], results[1])
 
     async def update_company_metrics(self, keys: Dict[str, str], event: str, quantity: int) -> None:
         """Update company metrics locally in cache (for track events)."""
@@ -893,10 +893,10 @@ def _clear_pending_requests(self) -> None:
                     fut.set_exception(RuntimeError("DataStream client disconnected"))
         self._pending_company.clear()
 
-        for futures in self._pending_user.values():
-            for fut in futures:
-                if not fut.done():
-                    fut.set_exception(RuntimeError("DataStream client disconnected"))
+        for user_futures in self._pending_user.values():
+            for user_fut in user_futures:
+                if not user_fut.done():
+                    user_fut.set_exception(RuntimeError("DataStream client disconnected"))
         self._pending_user.clear()
 
         if self._pending_flags is not None and not self._pending_flags.done():

src/schematic/datastream/websocket_client.py

@@ -112,7 +112,7 @@ def __init__(self, options: ClientOptions) -> None:
             raise ValueError("url is required")
         if not options.api_key:
             raise ValueError("api_key is required")
-        if not options.message_handler:
+        if options.message_handler is None:  # type: ignore[operator]
             raise ValueError("message_handler is required")
 
         # Auto-convert HTTP(S) URLs to WebSocket URLs

tests/datastream/test_merge.py

@@ -221,7 +221,7 @@ def test_sets_base_plan_to_none(self) -> None:
 class TestPartialCompanyMissingID:
     def test_raises_value_error(self) -> None:
         existing = base_company()
-        partial = {"traits": []}
+        partial: dict[str, list[str]] = {"traits": []}
 
         with pytest.raises(ValueError, match="missing required field: id"):
             partial_company(existing, partial)
@@ -314,6 +314,7 @@ def test_full_entity_partial_message(self) -> None:
         # Credit balances merge: existing credit-1 overwritten, credit-new added
         assert merged.credit_balances == {"credit-1": 999.0, "credit-new": 50.0}
 
+        assert merged.entitlements is not None
         assert len(merged.entitlements) == 2
         assert merged.entitlements[0].feature_id == "feat-new"
         assert merged.entitlements[1].feature_id == "feat-2"
@@ -493,13 +494,15 @@ def test_full_copy_is_independent(self) -> None:
         )
 
         cp = deep_copy_company(orig)
+        assert cp is not None
 
         assert cp.id == orig.id
         assert cp.account_id == orig.account_id
         assert cp.environment_id == orig.environment_id
         assert cp.base_plan_id == orig.base_plan_id
         assert cp.keys == orig.keys
         assert cp.credit_balances == orig.credit_balances
+        assert cp.subscription is not None
         assert cp.subscription.id == "sub-1"
         assert cp.metrics[0].value == 42
 
@@ -518,6 +521,7 @@ def test_empty_fields(self) -> None:
             rules=[],
         )
         cp = deep_copy_user(orig)
+        assert cp is not None
 
         assert cp.id == "u1"
         assert cp.keys == {}
@@ -528,6 +532,7 @@ def test_full_copy_is_independent(self) -> None:
         orig = base_user().model_copy(update={"rules": [_make_rule("r1")]})
 
         cp = deep_copy_user(orig)
+        assert cp is not None
 
         assert cp.id == orig.id
         assert cp.account_id == orig.account_id

tests/datastream/test_websocket_client.py

@@ -171,7 +171,7 @@ async def handler(msg): pass
 def test_init_missing_message_handler() -> None:
     with pytest.raises(ValueError, match="message_handler is required"):
         DatastreamWSClient(
-            ClientOptions(url="wss://example.com", api_key="key", message_handler=None, logger=logger)
+            ClientOptions(url="wss://example.com", api_key="key", message_handler=None, logger=logger)  # type: ignore[arg-type]
         )
 
 
@@ -236,7 +236,7 @@ async def test_string_message_delivered_to_handler() -> None:
 
 async def test_bytes_message_delivered_to_handler() -> None:
     payload = json.dumps({"entity_type": "rulesengine.Flag", "message_type": "full", "data": None})
-    ws = MockWebSocket(messages=[payload.encode()])
+    ws = MockWebSocket(messages=[payload.encode()])  # type: ignore[list-item]
     client, ws, received = make_client(ws=ws)
 
     with patch("schematic.datastream.websocket_client.websockets.connect", make_connect(ws)):
@@ -431,7 +431,7 @@ async def test_stops_at_max_reconnect_attempts() -> None:
         ClientOptions(
             url="wss://test.example.com/datastream",
             api_key="key",
-            message_handler=lambda _: None,
+            message_handler=lambda _: None,  # type: ignore[arg-type,return-value]
             logger=logger,
             min_reconnect_delay=0.0,
             max_reconnect_delay=0.0,