Fix pyright issues in Runloop rollout and related tests.

Use the correct blueprint build_context type and add narrowing assertions in new test coverage.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: tode-rl

examples/runloop_remote_rollout/create_blueprint.py

@@ -6,7 +6,6 @@
 from pathlib import Path
 
 from runloop_api_client import RunloopSDK
-from runloop_api_client.types.blueprint_build_parameters import BuildContext
 
 
 DEFAULT_DOCKERFILE = """\
@@ -60,7 +59,7 @@ def main() -> None:
     blueprint = runloop.blueprint.create(
         name=args.name,
         dockerfile=DEFAULT_DOCKERFILE,
-        build_context=BuildContext(type="object", object_id=build_context.id),
+        build_context={"type": "object", "object_id": build_context.id},
     )
 
     print(f"export RUNLOOP_BLUEPRINT_ID={blueprint.id}")

tests/adapters/test_fireworks_tracing_logprobs.py

@@ -71,6 +71,7 @@ def test_attaches_completion_logprobs_and_message_logprobs(self):
 
         assistant = row.messages[-1]
         assert assistant.role == "assistant"
+        assert assistant.logprobs is not None
         content = assistant.logprobs["content"]
         assert len(content) == len(extra["completion_logprobs"])
         assert content[0]["token_id"] == 10
@@ -83,10 +84,13 @@ def test_omits_token_id_keys_when_any_missing(self):
         assert row is not None
 
         extra = row.execution_metadata.extra
+        assert extra is not None
         assert "completion_logprobs" in extra
         assert "completion_token_ids" not in extra
 
-        content = row.messages[-1].logprobs["content"]
+        logprobs = row.messages[-1].logprobs
+        assert logprobs is not None
+        content = logprobs["content"]
         assert len(content) == 2
         assert all("token_id" not in entry for entry in content)
         assert content[0]["logprob"] == pytest.approx(-0.1)

tests/adapters/test_r3_deserializer.py

@@ -143,8 +143,9 @@ def test_all_mode_uint8(self):
         assert metadata["replayed_token_count"] == total_tokens
 
         for i in range(total_tokens):
-            assert matrices[i] is not None
-            decoded = base64.b64decode(matrices[i])
+            matrix = matrices[i]
+            assert matrix is not None
+            decoded = base64.b64decode(matrix)
             assert decoded == matrices_raw[i]
 
     def test_suffix_mode(self):
@@ -180,8 +181,9 @@ def test_suffix_mode(self):
         # Positions from start_token to start_token+replayed should have data
         for i in range(replayed):
             pos = start_token + i
-            assert matrices[pos] is not None
-            decoded = base64.b64decode(matrices[pos])
+            matrix = matrices[pos]
+            assert matrix is not None
+            decoded = base64.b64decode(matrix)
             assert decoded == matrices_raw[i]
 
     def test_bitmap_mode(self):
@@ -220,9 +222,10 @@ def test_bitmap_mode(self):
 
         for i in range(total_tokens):
             if i in replayed_positions:
-                assert matrices[i] is not None
+                matrix = matrices[i]
+                assert matrix is not None
                 idx = replayed_positions.index(i)
-                decoded = base64.b64decode(matrices[i])
+                decoded = base64.b64decode(matrix)
                 assert decoded == matrices_raw[idx]
             else:
                 assert matrices[i] is None
@@ -249,7 +252,9 @@ def test_uint16_dtype(self):
         assert metadata["routing_dtype"] == "uint16"
         assert len(matrices) == total_tokens
         for i in range(total_tokens):
-            decoded = base64.b64decode(matrices[i])
+            matrix = matrices[i]
+            assert matrix is not None
+            decoded = base64.b64decode(matrix)
             assert decoded == matrices_raw[i]
 
     def test_zero_replayed_tokens(self):
@@ -302,6 +307,7 @@ def test_high_compression_ratio_payload(self):
         assert len(matrices) == total_tokens
         assert metadata["replayed_token_count"] == total_tokens
         for m in matrices:
+            assert m is not None
             assert base64.b64decode(m) == b"\x00" * matrix_elem_size
 
 

tests/pytest/test_runloop_rollout_processor.py

@@ -1,6 +1,7 @@
 import asyncio
 from types import SimpleNamespace
 import urllib.error
+from email.message import Message
 
 import pytest
 
@@ -273,7 +274,9 @@ def read(self, size):
     def _urlopen(request, timeout):
         calls.append((request.full_url, timeout))
         if len(calls) == 1:
-            raise urllib.error.HTTPError(request.full_url, 503, "Service Unavailable", hdrs=None, fp=None)
+            raise urllib.error.HTTPError(
+                request.full_url, 503, "Service Unavailable", hdrs=Message(), fp=None
+            )
         return ReadyResponse()
 
     monkeypatch.setattr(runloop_rollout_processor_module.urllib.request, "urlopen", _urlopen)
@@ -296,7 +299,7 @@ def test_startup_wait_accepts_non_5xx_http_errors(monkeypatch):
 
     def _urlopen(request, timeout):
         calls.append((request.full_url, timeout))
-        raise urllib.error.HTTPError(request.full_url, 404, "Not Found", hdrs=None, fp=None)
+        raise urllib.error.HTTPError(request.full_url, 404, "Not Found", hdrs=Message(), fp=None)
 
     monkeypatch.setattr(runloop_rollout_processor_module.urllib.request, "urlopen", _urlopen)