From af4f0c9c2a4a22b8e988d178232bba6380653fb5 Mon Sep 17 00:00:00 2001 From: Joe Zou Date: Mon, 7 Oct 2024 15:46:34 -0700 Subject: [PATCH] Replace `mock` with `unittest.mock` in `model_maker` tests. PiperOrigin-RevId: 683353661 --- .../python/text/text_classifier/preprocessor_test.py | 3 +-- .../vision/gesture_recognizer/gesture_recognizer_test.py | 3 +-- .../python/vision/image_classifier/image_classifier_test.py | 3 +-- 3 files changed, 3 insertions(+), 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 index e830bfbf61..eea16e7a05 100644 --- 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 index 41799af975..d5e078e89d 100644 --- 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 index 21d01b70be..a96d08ece4 100644 --- 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(