Fix unit tests

openvinotoolkit · Apr 18, 2024 · 653ac02 · 653ac02
1 parent d1ba1f9
commit 653ac02
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/unit/core/model/test_visual_prompting.py b/tests/unit/core/model/test_visual_prompting.py
@@ -221,7 +221,7 @@ def test_optimization_config(self, otx_zero_shot_visual_prompting_model) -> None
 
     def test_on_test_start(self, mocker, otx_zero_shot_visual_prompting_model) -> None:
         """Test on_test_start."""
-        otx_zero_shot_visual_prompting_model.model.load_reference_info = Mock(return_value=False)
+        otx_zero_shot_visual_prompting_model.load_reference_info = Mock(return_value=False)
         otx_zero_shot_visual_prompting_model.trainer = Mock()
         mocker_run = mocker.patch.object(otx_zero_shot_visual_prompting_model.trainer.fit_loop, "run")
         mocker_setup_data = mocker.patch.object(
@@ -238,7 +238,7 @@ def test_on_test_start(self, mocker, otx_zero_shot_visual_prompting_model) -> No
 
     def test_on_predict_start(self, mocker, otx_zero_shot_visual_prompting_model) -> None:
         """Test on_predict_start."""
-        otx_zero_shot_visual_prompting_model.model.load_reference_info = Mock(return_value=False)
+        otx_zero_shot_visual_prompting_model.load_reference_info = Mock(return_value=False)
         otx_zero_shot_visual_prompting_model.trainer = Mock()
         mocker_run = mocker.patch.object(otx_zero_shot_visual_prompting_model.trainer.fit_loop, "run")
         mocker_setup_data = mocker.patch.object(
@@ -256,7 +256,7 @@ def test_on_predict_start(self, mocker, otx_zero_shot_visual_prompting_model) ->
     def test_on_train_epoch_end(self, mocker, tmpdir, otx_zero_shot_visual_prompting_model) -> None:
         """Test on_train_epoch_end."""
         otx_zero_shot_visual_prompting_model.save_outputs = True
-        otx_zero_shot_visual_prompting_model.model.save_reference_info = Mock()
+        otx_zero_shot_visual_prompting_model.save_reference_info = Mock()
         otx_zero_shot_visual_prompting_model.trainer = Mock()
         mocker.patch.object(otx_zero_shot_visual_prompting_model.trainer, "default_root_dir")