Patched os.environ in the unit test

googleapis · Mar 17, 2022 · 620043b · 620043b
1 parent bfe1e98
commit 620043b
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py
@@ -63,13 +63,18 @@ def mock_auth_default():
         assert initializer.global_config.project == _TEST_PROJECT
 
     def test_infer_project_id(self):
-        def mock_get_project_id():
+        cloud_project_number = 123
+
+        def mock_get_project_id(project_number: str, **_):
+            assert project_number == cloud_project_number
             return _TEST_PROJECT
 
         with mock.patch.object(
             target=resource_manager_utils,
             attribute="get_project_id",
             new=mock_get_project_id,
+        ), mock.patch.dict(
+            os.environ, {"CLOUD_ML_PROJECT_ID": cloud_project_number}, clear=True
         ):
             assert initializer.global_config.project == _TEST_PROJECT