Replace mock with unittest.mock in model_maker tests.

PiperOrigin-RevId: 683353661
google-ai-edge · Oct 7, 2024 · af4f0c9 · af4f0c9
1 parent eafcf31
commit af4f0c9
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 6 deletions.
diff --git a/mediapipe/model_maker/python/text/text_classifier/preprocessor_test.py b/mediapipe/model_maker/python/text/text_classifier/preprocessor_test.py
@@ -19,7 +19,6 @@
 from unittest import mock as unittest_mock
 
 from absl.testing import parameterized
-import mock
 import numpy as np
 import numpy.testing as npt
 import tensorflow as tf
@@ -158,7 +157,7 @@ def test_bert_preprocessor_cache(self):
 
     # The second time running preprocessor, it should load from cache directly
     mock_stdout = io.StringIO()
-    with mock.patch('sys.stdout', mock_stdout):
+    with unittest_mock.patch('sys.stdout', mock_stdout):
       _ = bert_preprocessor.preprocess(dataset)
     self.assertEqual(
         mock_stdout.getvalue(),

diff --git a/mediapipe/model_maker/python/vision/gesture_recognizer/gesture_recognizer_test.py b/mediapipe/model_maker/python/vision/gesture_recognizer/gesture_recognizer_test.py
@@ -18,7 +18,6 @@
 from unittest import mock as unittest_mock
 import zipfile
 
-import mock
 import tensorflow as tf
 
 from mediapipe.model_maker.python.core.utils import test_util
@@ -170,7 +169,7 @@ def test_continual_training_by_loading_checkpoint(self):
     gesture_recognizer_options = gesture_recognizer.GestureRecognizerOptions(
         model_options=mo, hparams=hparams)
     mock_stdout = io.StringIO()
-    with mock.patch('sys.stdout', mock_stdout):
+    with unittest_mock.patch('sys.stdout', mock_stdout):
       model = gesture_recognizer.GestureRecognizer.create(
           train_data=self._train_data,
           validation_data=self._validation_data,

diff --git a/mediapipe/model_maker/python/vision/image_classifier/image_classifier_test.py b/mediapipe/model_maker/python/vision/image_classifier/image_classifier_test.py
@@ -19,7 +19,6 @@
 
 from unittest import mock as unittest_mock
 from absl.testing import parameterized
-import mock
 import numpy as np
 import tensorflow as tf
 
@@ -143,7 +142,7 @@ def test_create_and_train_model(
 
   def test_continual_training_by_loading_checkpoint(self):
     mock_stdout = io.StringIO()
-    with mock.patch('sys.stdout', mock_stdout):
+    with unittest_mock.patch('sys.stdout', mock_stdout):
       options = image_classifier.ImageClassifierOptions(
           supported_model=image_classifier.SupportedModels.EFFICIENTNET_LITE0,
           hparams=image_classifier.HParams(