Add test case to model transform tests (#1795)

add test case
openvinotoolkit · Feb 29, 2024 · 83b53d8 · 83b53d8
1 parent 2e99ecd
commit 83b53d8
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/tests/unit/engine/test_setup_transform.py b/tests/unit/engine/test_setup_transform.py
@@ -219,3 +219,12 @@ def test_precendence_datamodule(self, checkpoint_path: Path) -> None:
         model = DummyModel()
         Engine._setup_transform(model, ckpt_path=checkpoint_path, datamodule=datamodule)  # noqa: SLF001
         assert model.transform == transform
+
+    def test_transform_already_assigned(self) -> None:
+        """Tests if the transform from the model is used when the model already has a transform assigned."""
+        transform = Transform()
+        model = DummyModel()
+        model.set_transform(transform)
+        datamodule = DummyDataModule()
+        Engine._setup_transform(model, datamodule=datamodule)  # noqa: SLF001
+        assert model.transform == transform