fix: address broken unit tests in certain environments

ucdmkt · copybara-github · commit d06b22d1ac61 · 2023-01-13T10:08:46.000-08:00
PiperOrigin-RevId: 501875885
diff --git a/google/cloud/aiplatform/vizier/pyvizier/study_config.py b/google/cloud/aiplatform/vizier/pyvizier/study_config.py
@@ -117,19 +117,36 @@ class SearchSpace(SearchSpace):
     @classmethod
     def from_proto(cls, proto: study_pb2.StudySpec) -> "SearchSpace":
         """Extracts a SearchSpace object from a StudyConfig proto."""
-        parameter_configs = []
+
+        # For google-vizier <= 0.0.15
+        if hasattr(cls, "_factory"):
+            parameter_configs = []
+            for pc in proto.parameters:
+                parameter_configs.append(
+                    proto_converters.ParameterConfigConverter.from_proto(pc)
+                )
+            return cls._factory(parameter_configs=parameter_configs)
+
+        result = cls()
         for pc in proto.parameters:
-            parameter_configs.append(
-                proto_converters.ParameterConfigConverter.from_proto(pc)
-            )
-        return cls._factory(parameter_configs=parameter_configs)
+            result.add(proto_converters.ParameterConfigConverter.from_proto(pc))
+
+        return result
 
     @property
     def parameter_protos(self) -> List[study_pb2.StudySpec.ParameterSpec]:
         """Returns the search space as a List of ParameterConfig protos."""
+
+        # For google-vizier <= 0.0.15
+        if isinstance(self._parameter_configs, list):
+            return [
+                proto_converters.ParameterConfigConverter.to_proto(pc)
+                for pc in self._parameter_configs
+            ]
+
         return [
             proto_converters.ParameterConfigConverter.to_proto(pc)
-            for pc in self._parameter_configs
+            for _, pc in self._parameter_configs.items()
         ]
 
 
diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py
@@ -614,6 +614,7 @@ def list_artifact_empty_mock():
         yield list_artifacts_mock
 
 
+@pytest.mark.usefixtures("google_auth_mock")
 class TestExecution:
     def setup_method(self):
         reload(initializer)
@@ -893,6 +894,7 @@ def test_query_input_and_output_artifacts(
         assert artifact_list[0]._gca_resource == expected_artifact
 
 
+@pytest.mark.usefixtures("google_auth_mock")
 class TestArtifact:
     def setup_method(self):
         reload(initializer)
diff --git a/tests/unit/aiplatform/test_metadata_store.py b/tests/unit/aiplatform/test_metadata_store.py
@@ -134,6 +134,7 @@ def delete_metadata_store_mock():
         yield delete_metadata_store_mock
 
 
+@pytest.mark.usefixtures("google_auth_mock")
 class TestMetadataStore:
     def setup_method(self):
         reload(initializer)
diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py
@@ -390,6 +390,7 @@ def test_wrapped_client():
     )
 
 
+@pytest.mark.usefixtures("google_auth_mock")
 def test_client_w_override_default_version():
 
     test_client_info = gapic_v1.client_info.ClientInfo()
@@ -407,6 +408,7 @@ def test_client_w_override_default_version():
     )
 
 
+@pytest.mark.usefixtures("google_auth_mock")
 def test_client_w_override_select_version():
 
     test_client_info = gapic_v1.client_info.ClientInfo()

Original file line number	Diff line number	Diff line change
`@@ -390,6 +390,7 @@ def test_wrapped_client():`
`390`	`390`	`)`
`391`	`391`
`392`	`392`
	`393`	`+@pytest.mark.usefixtures("google_auth_mock")`
`393`	`394`	`def test_client_w_override_default_version():`
`394`	`395`
`395`	`396`	`test_client_info = gapic_v1.client_info.ClientInfo()`
`@@ -407,6 +408,7 @@ def test_client_w_override_default_version():`
`407`	`408`	`)`
`408`	`409`
`409`	`410`
	`411`	`+@pytest.mark.usefixtures("google_auth_mock")`
`410`	`412`	`def test_client_w_override_select_version():`
`411`	`413`
`412`	`414`	`test_client_info = gapic_v1.client_info.ClientInfo()`