Typing in tests: added asserts and cast to remove some mypy errors

test/base/test_batch.py

@@ -2,6 +2,7 @@
 import pickle
 import sys
 from itertools import starmap
+from typing import cast
 
 import networkx as nx
 import numpy as np
@@ -160,7 +161,11 @@ def test_batch() -> None:
     batch5 = Batch(a=np.array([{"index": 0}]))
     assert isinstance(batch5.a, Batch)
     assert np.allclose(batch5.a.index, [0])
+    # We use setattr b/c the setattr of Batch will actually change the type of the field that is being set!
+    # However, mypy would not understand this, and rightly expect that batch.b = some_array would lead to
+    # batch.b being an array (which it is not, it's turned into a Batch instead)
     batch5.b = np.array([{"index": 1}])
+    batch5.b = cast(Batch, batch5.b)
     assert isinstance(batch5.b, Batch)
     assert np.allclose(batch5.b.index, [1])
 

test/base/test_collector.py

@@ -314,9 +314,11 @@ def test_collector_with_dict_state() -> None:
     batch, _ = c1.buffer.sample(10)
     c0.buffer.update(c1.buffer)
     assert len(c0.buffer) in [42, 43]
+    cur_obs = c0.buffer[:].obs
+    assert isinstance(cur_obs, Batch)
     if len(c0.buffer) == 42:
         assert np.all(
-            c0.buffer[:].obs.index[..., 0]
+            cur_obs.index[..., 0]
             == [
                 0,
                 1,
@@ -364,7 +366,7 @@ def test_collector_with_dict_state() -> None:
         ), c0.buffer[:].obs.index[..., 0]
     else:
         assert np.all(
-            c0.buffer[:].obs.index[..., 0]
+            cur_obs.index[..., 0]
             == [
                 0,
                 1,