From 010e94d731fe3a90a048cb0c4fdc31a8f7ec5648 Mon Sep 17 00:00:00 2001 From: zhitaoli Date: Tue, 1 Feb 2022 14:31:02 -0800 Subject: [PATCH] Replaced deprecated assertDictContainsSubset with assertLessEqual(itemsA, itemsB) PiperOrigin-RevId: 425717602 --- RELEASE.md | 5 +- .../local_docker_runner_test.py | 20 ++--- .../google_cloud_ai_platform/runner_test.py | 74 +++++++++---------- 3 files changed, 50 insertions(+), 49 deletions(-) diff --git a/RELEASE.md b/RELEASE.md index e358448f16..09c11f0794 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -12,7 +12,10 @@ ## Bug Fixes and Other Changes -* Fixed the cluster spec error in CAIP Tuner on Vertex when `num_parallel_trials = 1` +* Fixed the cluster spec error in CAIP Tuner on Vertex when + `num_parallel_trials = 1` +* Replaced deprecated assertDictContainsSubset with + assertLessEqual(itemsA, itemsB). ## Dependency Updates diff --git a/tfx/components/infra_validator/model_server_runners/local_docker_runner_test.py b/tfx/components/infra_validator/model_server_runners/local_docker_runner_test.py index 9444497863..8b7a0b0fe8 100644 --- a/tfx/components/infra_validator/model_server_runners/local_docker_runner_test.py +++ b/tfx/components/infra_validator/model_server_runners/local_docker_runner_test.py @@ -84,16 +84,16 @@ def testStart(self): # Check calls. self._docker_client.containers.run.assert_called() _, run_kwargs = self._docker_client.containers.run.call_args - self.assertDictContainsSubset(dict( - image='tensorflow/serving:1.15.0', - environment={ - 'MODEL_NAME': 'chicago-taxi', - 'MODEL_BASE_PATH': '/model' - }, - publish_all_ports=True, - auto_remove=True, - detach=True - ), run_kwargs) + self.assertLessEqual( + dict( + image='tensorflow/serving:1.15.0', + environment={ + 'MODEL_NAME': 'chicago-taxi', + 'MODEL_BASE_PATH': '/model' + }, + publish_all_ports=True, + auto_remove=True, + detach=True).items(), run_kwargs.items()) def testStartMultipleTimesFail(self): # Prepare mocks and variables. diff --git a/tfx/extensions/google_cloud_ai_platform/runner_test.py b/tfx/extensions/google_cloud_ai_platform/runner_test.py index 95fe8b27d7..8c0089da7b 100644 --- a/tfx/extensions/google_cloud_ai_platform/runner_test.py +++ b/tfx/extensions/google_cloud_ai_platform/runner_test.py @@ -114,21 +114,20 @@ def testStartCloudTraining(self, mock_discovery): default_image = 'gcr.io/tfx-oss-public/tfx:{}'.format( version_utils.get_image_version()) - self.assertDictContainsSubset( - { - 'masterConfig': { - 'imageUri': - default_image, - 'containerCommand': - runner._CONTAINER_COMMAND + [ - '--executor_class_path', class_path, '--inputs', '{}', - '--outputs', '{}', '--exec-properties', - ('{"custom_config": ' - '"{\\"ai_platform_training_args\\": {\\"project\\": \\"12345\\"' - '}}"}') - ], - }, - }, body['training_input']) + self.assertLessEqual({ + 'masterConfig': { + 'imageUri': + default_image, + 'containerCommand': + runner._CONTAINER_COMMAND + [ + '--executor_class_path', class_path, '--inputs', '{}', + '--outputs', '{}', '--exec-properties', + ('{"custom_config": ' + '"{\\"ai_platform_training_args\\": {\\"project\\": \\"12345\\"' + '}}"}') + ], + }, + }.items(), body['training_input'].items()) self.assertNotIn('project', body['training_input']) self.assertStartsWith(body['job_id'], 'tfx_') self._mock_get.execute.assert_called_with() @@ -239,28 +238,27 @@ def testStartCloudTrainingWithUserContainer_Vertex(self, mock_gapic): custom_job=mock.ANY) kwargs = self._mock_create.call_args[1] body = kwargs['custom_job'] - self.assertDictContainsSubset( - { - 'worker_pool_specs': [{ - 'container_spec': { - 'image_uri': - 'my-custom-image', - 'command': - runner._CONTAINER_COMMAND + [ - '--executor_class_path', class_path, '--inputs', - '{}', '--outputs', '{}', '--exec-properties', - ('{"custom_config": ' - '"{\\"ai_platform_training_args\\": ' - '{\\"project\\": \\"12345\\", ' - '\\"worker_pool_specs\\": ' - '[{\\"container_spec\\": ' - '{\\"image_uri\\": \\"my-custom-image\\"}}]}, ' - '\\"ai_platform_training_job_id\\": ' - '\\"my_jobid\\"}"}') - ], - }, - },], - }, body['job_spec']) + self.assertLessEqual({ + 'worker_pool_specs': [{ + 'container_spec': { + 'image_uri': + 'my-custom-image', + 'command': + runner._CONTAINER_COMMAND + [ + '--executor_class_path', class_path, '--inputs', '{}', + '--outputs', '{}', '--exec-properties', + ('{"custom_config": ' + '"{\\"ai_platform_training_args\\": ' + '{\\"project\\": \\"12345\\", ' + '\\"worker_pool_specs\\": ' + '[{\\"container_spec\\": ' + '{\\"image_uri\\": \\"my-custom-image\\"}}]}, ' + '\\"ai_platform_training_job_id\\": ' + '\\"my_jobid\\"}"}') + ], + }, + },], + }.items(), body['job_spec'].items()) self.assertEqual(body['display_name'], 'my_jobid') self._mock_get.assert_called_with(name='vertex_job_study_id') @@ -329,7 +327,7 @@ def testStartCloudTrainingWithVertexCustomJob(self, mock_gapic): }, body['job_spec']) self.assertEqual(body['display_name'], 'valid_name') self.assertDictEqual(body['encryption_spec'], expected_encryption_spec) - self.assertDictContainsSubset(user_provided_labels, body['labels']) + self.assertLessEqual(user_provided_labels.items(), body['labels'].items()) self._mock_get.assert_called_with(name='vertex_job_study_id') def _setUpPredictionMocks(self):