split tests

Author: Borda

tests/tests_pytorch/trainer/connectors/test_callback_connector.py

@@ -104,12 +104,12 @@ def test_checkpoint_callbacks_are_last(tmp_path):
     ]
 
 
-class StatefulCallback0(Callback):
+class StatefulCallbackContent0(Callback):
     def state_dict(self):
         return {"content0": 0}
 
 
-class StatefulCallback1(Callback):
+class StatefulCallbackContent1(Callback):
     def __init__(self, unique=None, other=None):
         self._unique = unique
         self._other = other
@@ -126,9 +126,9 @@ def test_all_callback_states_saved_before_checkpoint_callback(tmp_path):
     """Test that all callback states get saved even if the ModelCheckpoint is not given as last and when there are
     multiple callbacks of the same type."""
 
-    callback0 = StatefulCallback0()
-    callback1 = StatefulCallback1(unique="one")
-    callback2 = StatefulCallback1(unique="two", other=2)
+    callback0 = StatefulCallbackContent0()
+    callback1 = StatefulCallbackContent1(unique="one")
+    callback2 = StatefulCallbackContent1(unique="two", other=2)
     checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, filename="all_states")
     model = BoringModel()
     trainer = Trainer(
@@ -147,9 +147,9 @@ def test_all_callback_states_saved_before_checkpoint_callback(tmp_path):
     trainer.fit(model)
 
     ckpt = torch.load(str(tmp_path / "all_states.ckpt"), weights_only=True)
-    state0 = ckpt["callbacks"]["StatefulCallback0"]
-    state1 = ckpt["callbacks"]["StatefulCallback1{'unique': 'one'}"]
-    state2 = ckpt["callbacks"]["StatefulCallback1{'unique': 'two'}"]
+    state0 = ckpt["callbacks"]["StatefulCallbackContent0"]
+    state1 = ckpt["callbacks"]["StatefulCallbackContent1{'unique': 'one'}"]
+    state2 = ckpt["callbacks"]["StatefulCallbackContent1{'unique': 'two'}"]
     assert "content0" in state0
     assert state0["content0"] == 0
     assert "content1" in state1
@@ -325,72 +325,65 @@ def state_dict(self):
         Trainer(callbacks=[MockCallback(), MockCallback()])
 
 
-def test_validate_callbacks_list_function():
-    """Test the _validate_callbacks_list function directly with various scenarios."""
-
-    # Test with non-stateful callbacks
-    callback1 = Callback()
-    callback2 = Callback()
-    _validate_callbacks_list([callback1, callback2])
-
-    # Test with single stateful callback
-    class StatefulCallback(Callback):
-        def state_dict(self):
-            return {"state": 1}
-
-    stateful_cb = StatefulCallback()
-    _validate_callbacks_list([stateful_cb])
-
-    # Test with multiple stateful callbacks with unique state keys
-    class StatefulCallback1(Callback):
-        @property
-        def state_key(self):
-            return "unique_key_1"
-
-        def state_dict(self):
-            return {"state": 1}
-
-    class StatefulCallback2(Callback):
-        @property
-        def state_key(self):
-            return "unique_key_2"
-
-        def state_dict(self):
-            return {"state": 2}
-
-    stateful_cb1 = StatefulCallback1()
-    stateful_cb2 = StatefulCallback2()
-    _validate_callbacks_list([stateful_cb1, stateful_cb2])
+# Test with single stateful callback
+class StatefulCallback(Callback):
+    def state_dict(self):
+        return {"state": 1}
 
-    # Test with multiple stateful callbacks with same state key
-    class ConflictingCallback(Callback):
-        @property
-        def state_key(self):
-            return "same_key"
+# Test with multiple stateful callbacks with unique state keys
+class StatefulCallback1(Callback):
+    @property
+    def state_key(self):
+        return "unique_key_1"
 
-        def state_dict(self):
-            return {"state": 1}
+    def state_dict(self):
+        return {"state": 1}
 
-    conflicting_cb1 = ConflictingCallback()
-    conflicting_cb2 = ConflictingCallback()
+class StatefulCallback2(Callback):
+    @property
+    def state_key(self):
+        return "unique_key_2"
 
-    with pytest.raises(RuntimeError, match="Found more than one stateful callback of type `ConflictingCallback`"):
-        _validate_callbacks_list([conflicting_cb1, conflicting_cb2])
+    def state_dict(self):
+        return {"state": 2}
+
+@pytest.mark.parametrize(
+    ("callbacks"),
+    [[Callback(), Callback()],
+     [StatefulCallback()],
+     [StatefulCallback1(), StatefulCallback2()],
+     ],
+)
+def test_validate_callbacks_list_function(callbacks: list):
+    """Test the _validate_callbacks_list function directly with various scenarios."""
+    _validate_callbacks_list(callbacks)
 
-    # Test with mix of stateful and non-stateful callbacks where stateful ones conflict
-    with pytest.raises(RuntimeError, match="Found more than one stateful callback of type `ConflictingCallback`"):
-        _validate_callbacks_list([callback1, conflicting_cb1, callback2, conflicting_cb2])
 
-    # Test with different types of stateful callbacks that happen to have same state key
-    class AnotherConflictingCallback(Callback):
-        @property
-        def state_key(self):
-            return "same_key"  # Same key as ConflictingCallback
+# Test with multiple stateful callbacks with same state key
+class ConflictingCallback(Callback):
+    @property
+    def state_key(self):
+        return "same_key"
 
-        def state_dict(self):
-            return {"state": 3}
+    def state_dict(self):
+        return {"state": 1}
 
-    another_conflicting_cb = AnotherConflictingCallback()
+# Test with different types of stateful callbacks that happen to have same state key
+class AnotherConflictingCallback(Callback):
+    @property
+    def state_key(self):
+        return "same_key"  # Same key as ConflictingCallback
 
-    with pytest.raises(RuntimeError, match="Found more than one stateful callback"):
-        _validate_callbacks_list([conflicting_cb1, another_conflicting_cb])
+    def state_dict(self):
+        return {"state": 3}
+@pytest.mark.parametrize(
+    ("callbacks", "match_msg"),
+    [
+        ([ConflictingCallback(), ConflictingCallback()], "Found more than one stateful callback of type `ConflictingCallback`"),
+        ([Callback(), ConflictingCallback(), Callback(),ConflictingCallback(), ], "Found more than one stateful callback of type `ConflictingCallback`"),
+        ([ConflictingCallback(), AnotherConflictingCallback()], "Found more than one stateful callback"),
+    ],
+)
+def test_raising_error_validate_callbacks_list_function(callbacks: list, match_msg:str):
+    with pytest.raises(RuntimeError, match=match_msg):
+        _validate_callbacks_list(callbacks)