diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py index adf547e540..19b5f7d468 100644 --- a/google/cloud/aiplatform/base.py +++ b/google/cloud/aiplatform/base.py @@ -727,7 +727,7 @@ def __repr__(self) -> str: def to_dict(self) -> Dict[str, Any]: """Returns the resource proto as a dictionary.""" - return json_format.MessageToDict(self.gca_resource._pb) + return json_format.MessageToDict(self._gca_resource._pb) @classmethod def _generate_display_name(cls, prefix: Optional[str] = None) -> str: diff --git a/google/cloud/aiplatform/metadata/artifact.py b/google/cloud/aiplatform/metadata/artifact.py index d01e9c9eb6..5f4f94e18b 100644 --- a/google/cloud/aiplatform/metadata/artifact.py +++ b/google/cloud/aiplatform/metadata/artifact.py @@ -31,7 +31,6 @@ from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform.metadata import resource from google.cloud.aiplatform.metadata import utils as metadata_utils -from google.cloud.aiplatform.metadata.schema import base_artifact from google.cloud.aiplatform.utils import rest_utils @@ -327,60 +326,15 @@ def create( credentials=credentials, ) - @classmethod - def create_from_base_artifact_schema( - cls, - *, - base_artifact_schema: "base_artifact.BaseArtifactSchema", - metadata_store_id: Optional[str] = "default", - project: Optional[str] = None, - location: Optional[str] = None, - credentials: Optional[auth_credentials.Credentials] = None, - ) -> "Artifact": - """Creates a new Metadata Artifact from a BaseArtifactSchema class instance. - - Args: - base_artifact_schema (BaseArtifactSchema): - Required. An instance of the BaseArtifactType class that can be - provided instead of providing artifact specific parameters. - metadata_store_id (str): - Optional. The portion of the resource name with - the format: - projects/123/locations/us-central1/metadataStores//artifacts/ - If not provided, the MetadataStore's ID will be set to "default". - project (str): - Optional. Project used to create this Artifact. Overrides project set in - aiplatform.init. - location (str): - Optional. Location used to create this Artifact. Overrides location set in - aiplatform.init. - credentials (auth_credentials.Credentials): - Optional. Custom credentials used to create this Artifact. Overrides - credentials set in aiplatform.init. - - Returns: - Artifact: Instantiated representation of the managed Metadata Artifact. - """ - - return cls.create( - resource_id=base_artifact_schema.artifact_id, - schema_title=base_artifact_schema.schema_title, - uri=base_artifact_schema.uri, - display_name=base_artifact_schema.display_name, - schema_version=base_artifact_schema.schema_version, - description=base_artifact_schema.description, - metadata=base_artifact_schema.metadata, - state=base_artifact_schema.state, - metadata_store_id=metadata_store_id, - project=project, - location=location, - credentials=credentials, - ) - @property def uri(self) -> Optional[str]: "Uri for this Artifact." - return self.gca_resource.uri + return self._gca_resource.uri + + @property + def state(self) -> Optional[gca_artifact.Artifact.State]: + "The State for this Artifact." + return self._gca_resource.state @classmethod def get_with_uri( diff --git a/google/cloud/aiplatform/metadata/execution.py b/google/cloud/aiplatform/metadata/execution.py index 895417fc64..9a85bce36f 100644 --- a/google/cloud/aiplatform/metadata/execution.py +++ b/google/cloud/aiplatform/metadata/execution.py @@ -31,7 +31,6 @@ from google.cloud.aiplatform.metadata import artifact from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform.metadata import resource -from google.cloud.aiplatform.metadata.schema import base_execution class Execution(resource._Resource): @@ -167,57 +166,6 @@ def create( return self - @classmethod - def create_from_base_execution_schema( - cls, - *, - base_execution_schema: "base_execution.BaseExecutionSchema", - metadata_store_id: Optional[str] = "default", - project: Optional[str] = None, - location: Optional[str] = None, - credentials: Optional[auth_credentials.Credentials] = None, - ) -> "Execution": - """ - Creates a new Metadata Execution. - - Args: - base_execution_schema (BaseExecutionSchema): - An instance of the BaseExecutionSchema class that can be - provided instead of providing schema specific parameters. - metadata_store_id (str): - Optional. The portion of the resource name with - the format: - projects/123/locations/us-central1/metadataStores//artifacts/ - If not provided, the MetadataStore's ID will be set to "default". - project (str): - Optional. Project used to create this Execution. Overrides project set in - aiplatform.init. - location (str): - Optional. Location used to create this Execution. Overrides location set in - aiplatform.init. - credentials (auth_credentials.Credentials): - Optional. Custom credentials used to create this Execution. Overrides - credentials set in aiplatform.init. - - Returns: - Execution: Instantiated representation of the managed Metadata Execution. - - """ - resource = Execution.create( - state=base_execution_schema.state, - schema_title=base_execution_schema.schema_title, - resource_id=base_execution_schema.execution_id, - display_name=base_execution_schema.display_name, - schema_version=base_execution_schema.schema_version, - metadata=base_execution_schema.metadata, - description=base_execution_schema.description, - metadata_store_id=metadata_store_id, - project=project, - location=location, - credentials=credentials, - ) - return resource - def __enter__(self): if self.state is not gca_execution.Execution.State.RUNNING: self.update(state=gca_execution.Execution.State.RUNNING) diff --git a/google/cloud/aiplatform/metadata/resource.py b/google/cloud/aiplatform/metadata/resource.py index fccd0b18d6..00133ba789 100644 --- a/google/cloud/aiplatform/metadata/resource.py +++ b/google/cloud/aiplatform/metadata/resource.py @@ -114,6 +114,14 @@ def schema_title(self) -> str: def description(self) -> str: return self._gca_resource.description + @property + def display_name(self) -> str: + return self._gca_resource.display_name + + @property + def schema_version(self) -> str: + return self._gca_resource.schema_version + @classmethod def get_or_create( cls, @@ -221,7 +229,7 @@ def get( Returns: resource (_Resource): - Instantiated representation of the managed Metadata resource or None if no resouce was found. + Instantiated representation of the managed Metadata resource or None if no resource was found. """ resource = cls._get( diff --git a/google/cloud/aiplatform/metadata/schema/base_artifact.py b/google/cloud/aiplatform/metadata/schema/base_artifact.py index c89d989edd..357ce9d58a 100644 --- a/google/cloud/aiplatform/metadata/schema/base_artifact.py +++ b/google/cloud/aiplatform/metadata/schema/base_artifact.py @@ -26,7 +26,7 @@ from google.cloud.aiplatform.metadata import constants -class BaseArtifactSchema(metaclass=abc.ABCMeta): +class BaseArtifactSchema(artifact.Artifact): """Base class for Metadata Artifact types.""" @property @@ -81,13 +81,40 @@ def __init__( Pipelines), and the system does not prescribe or check the validity of state transitions. """ + # resource_id is not stored in the proto. Create method uses the + # resource_id along with project_id and location to construct an + # resource_name which is stored in the proto message. self.artifact_id = artifact_id - self.uri = uri - self.display_name = display_name - self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION - self.description = description - self.metadata = metadata - self.state = state + + # Store all other attributes using the proto structure. + self._gca_resource = gca_artifact.Artifact() + self._gca_resource.uri = uri + self._gca_resource.display_name = display_name + self._gca_resource.schema_version = ( + schema_version or constants._DEFAULT_SCHEMA_VERSION + ) + self._gca_resource.description = description + + # If metadata is None covert to {} + metadata = metadata if metadata else {} + self._nested_update_metadata(self._gca_resource, metadata) + self._gca_resource.state = state + + # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8 + def _init_with_resource_name( + self, + *, + artifact_name: str, + ): + + """Initializes the Artifact instance using an existing resource. + + Args: + artifact_name (str): + Artifact name with the following format, this is globally unique in a metadataStore: + projects/123/locations/us-central1/metadataStores//artifacts/. + """ + super(BaseArtifactSchema, self).__init__(artifact_name=artifact_name) def create( self, @@ -117,10 +144,27 @@ def create( Returns: Artifact: Instantiated representation of the managed Metadata Artifact. """ - return artifact.Artifact.create_from_base_artifact_schema( - base_artifact_schema=self, + + # Check if metadata exists to avoid proto read error + metadata = None + if self._gca_resource.metadata: + metadata = self.metadata + + new_artifact_instance = artifact.Artifact.create( + resource_id=self.artifact_id, + schema_title=self.schema_title, + uri=self.uri, + display_name=self.display_name, + schema_version=self.schema_version, + description=self.description, + metadata=metadata, + state=self.state, metadata_store_id=metadata_store_id, project=project, location=location, credentials=credentials, ) + + # Reinstantiate this class using the newly created resource. + self._init_with_resource_name(artifact_name=new_artifact_instance.resource_name) + return self diff --git a/google/cloud/aiplatform/metadata/schema/base_context.py b/google/cloud/aiplatform/metadata/schema/base_context.py index a39835da00..b6d7f5b4d7 100644 --- a/google/cloud/aiplatform/metadata/schema/base_context.py +++ b/google/cloud/aiplatform/metadata/schema/base_context.py @@ -21,11 +21,12 @@ from google.auth import credentials as auth_credentials +from google.cloud.aiplatform.compat.types import context as gca_context from google.cloud.aiplatform.metadata import constants from google.cloud.aiplatform.metadata import context -class BaseContextSchema(metaclass=abc.ABCMeta): +class BaseContextSchema(context.Context): """Base class for Metadata Context schema.""" @property @@ -62,11 +63,35 @@ def __init__( description (str): Optional. Describes the purpose of the Context to be created. """ + # resource_id is not stored in the proto. Create method uses the + # resource_id along with project_id and location to construct an + # resource_name which is stored in the proto message. self.context_id = context_id - self.display_name = display_name - self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION - self.metadata = metadata - self.description = description + + # Store all other attributes using the proto structure. + self._gca_resource = gca_context.Context() + self._gca_resource.display_name = display_name + self._gca_resource.schema_version = ( + schema_version or constants._DEFAULT_SCHEMA_VERSION + ) + # If metadata is None covert to {} + metadata = metadata if metadata else {} + self._nested_update_metadata(self._gca_resource, metadata) + self._gca_resource.description = description + + # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8 + def _init_with_resource_name( + self, + *, + context_name: str, + ): + """Initializes the Artifact instance using an existing resource. + Args: + context_name (str): + Context name with the following format, this is globally unique in a metadataStore: + projects/123/locations/us-central1/metadataStores//contexts/. + """ + super(BaseContextSchema, self).__init__(resource_name=context_name) def create( self, @@ -97,15 +122,24 @@ def create( Context: Instantiated representation of the managed Metadata Context. """ - return context.Context.create( + # Check if metadata exists to avoid proto read error + metadata = None + if self._gca_resource.metadata: + metadata = self.metadata + + new_context = context.Context.create( resource_id=self.context_id, schema_title=self.schema_title, display_name=self.display_name, schema_version=self.schema_version, description=self.description, - metadata=self.metadata, + metadata=metadata, metadata_store_id=metadata_store_id, project=project, location=location, credentials=credentials, ) + + # Reinstantiate this class using the newly created resource. + self._init_with_resource_name(context_name=new_context.resource_name) + return self diff --git a/google/cloud/aiplatform/metadata/schema/base_execution.py b/google/cloud/aiplatform/metadata/schema/base_execution.py index c8d55a0637..1cbf66b825 100644 --- a/google/cloud/aiplatform/metadata/schema/base_execution.py +++ b/google/cloud/aiplatform/metadata/schema/base_execution.py @@ -27,7 +27,7 @@ from google.cloud.aiplatform.metadata import metadata -class BaseExecutionSchema(metaclass=abc.ABCMeta): +class BaseExecutionSchema(execution.Execution): """Base class for Metadata Execution schema.""" @property @@ -69,12 +69,38 @@ def __init__( description (str): Optional. Describes the purpose of the Execution to be created. """ - self.state = state + + # resource_id is not stored in the proto. Create method uses the + # resource_id along with project_id and location to construct an + # resource_name which is stored in the proto message. self.execution_id = execution_id - self.display_name = display_name - self.schema_version = schema_version or constants._DEFAULT_SCHEMA_VERSION - self.metadata = metadata - self.description = description + + # Store all other attributes using the proto structure. + self._gca_resource = gca_execution.Execution() + self._gca_resource.state = state + self._gca_resource.display_name = display_name + self._gca_resource.schema_version = ( + schema_version or constants._DEFAULT_SCHEMA_VERSION + ) + # If metadata is None covert to {} + metadata = metadata if metadata else {} + self._nested_update_metadata(self._gca_resource, metadata) + self._gca_resource.description = description + + # TODO() Switch to @singledispatchmethod constructor overload after py>=3.8 + def _init_with_resource_name( + self, + *, + execution_name: str, + ): + + """Initializes the Execution instance using an existing resource. + Args: + execution_name (str): + The Execution name with the following format, this is globally unique in a metadataStore. + projects/123/locations/us-central1/metadataStores//executions/. + """ + super(BaseExecutionSchema, self).__init__(execution_name=execution_name) def create( self, @@ -105,14 +131,29 @@ def create( Execution: Instantiated representation of the managed Metadata Execution. """ - self.execution = execution.Execution.create_from_base_execution_schema( - base_execution_schema=self, + # Check if metadata exists to avoid proto read error + metadata = None + if self._gca_resource.metadata: + metadata = self.metadata + + new_execution_instance = execution.Execution.create( + resource_id=self.execution_id, + schema_title=self.schema_title, + display_name=self.display_name, + schema_version=self.schema_version, + description=self.description, + metadata=metadata, + state=self.state, metadata_store_id=metadata_store_id, project=project, location=location, credentials=credentials, ) - return self.execution + # Reinstantiate this class using the newly created resource. + self._init_with_resource_name( + execution_name=new_execution_instance.resource_name + ) + return self def start_execution( self, @@ -172,7 +213,7 @@ def start_execution( f"metadata_store_id {metadata_store_id} is not supported. Only the default MetadataStore ID is supported." ) - return metadata._ExperimentTracker().start_execution( + new_execution_instance = metadata._ExperimentTracker().start_execution( schema_title=self.schema_title, display_name=self.display_name, resource_id=self.execution_id, @@ -185,3 +226,9 @@ def start_execution( location=location, credentials=credentials, ) + + # Reinstantiate this class using the newly created resource. + self._init_with_resource_name( + execution_name=new_execution_instance.resource_name + ) + return self diff --git a/google/cloud/aiplatform/metadata/schema/utils.py b/google/cloud/aiplatform/metadata/schema/utils.py index 72577d9324..c742824b76 100644 --- a/google/cloud/aiplatform/metadata/schema/utils.py +++ b/google/cloud/aiplatform/metadata/schema/utils.py @@ -159,7 +159,7 @@ def create_uri_from_resource_name(resource_name: str) -> bool: """ # TODO: support nested resource names such as models/123/evaluations/456 match_results = re.match( - r"^projects\/[A-Za-z0-9-]*\/locations\/([A-Za-z0-9-]*)\/[A-Za-z0-9-]*\/[A-Za-z0-9-]*$", + r"^projects\/[A-Za-z0-9-]*\/locations\/([A-Za-z0-9-]*)(\/metadataStores\/[A-Za-z0-9-]*)?(\/[A-Za-z0-9-]*\/[A-Za-z0-9-]*)+$", resource_name, ) if not match_results: diff --git a/samples/model-builder/conftest.py b/samples/model-builder/conftest.py index 3c5960a2b1..e312268b65 100644 --- a/samples/model-builder/conftest.py +++ b/samples/model-builder/conftest.py @@ -686,6 +686,24 @@ def mock_create_artifact(mock_artifact): yield mock_create_artifact +@pytest.fixture +def mock_create_schema_base_artifact(mock_artifact): + with patch.object( + aiplatform.metadata.schema.base_artifact.BaseArtifactSchema, "create" + ) as mock_create_schema_base_artifact: + mock_create_schema_base_artifact.return_value = mock_artifact + yield mock_create_schema_base_artifact + + +@pytest.fixture +def mock_create_schema_base_execution(mock_execution): + with patch.object( + aiplatform.metadata.schema.base_execution.BaseExecutionSchema, "create" + ) as mock_create_schema_base_execution: + mock_create_schema_base_execution.return_value = mock_execution + yield mock_create_schema_base_execution + + @pytest.fixture def mock_list_artifact(mock_artifact): with patch.object(aiplatform.Artifact, "list") as mock_list_artifact: diff --git a/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py b/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py index 09b9249ff2..ecd013516d 100644 --- a/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py +++ b/samples/model-builder/experiment_tracking/create_artifact_with_sdk_sample_test.py @@ -14,12 +14,10 @@ import create_artifact_with_sdk_sample -from google.cloud.aiplatform.compat.types import artifact as gca_artifact - import test_constants as constants -def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_artifact): +def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_schema_base_artifact): artifact = create_artifact_with_sdk_sample.create_artifact_sample( project=constants.PROJECT, location=constants.LOCATION, @@ -31,19 +29,7 @@ def test_create_artifact_with_sdk_sample(mock_artifact, mock_create_artifact): metadata=constants.METADATA, ) - mock_create_artifact.assert_called_with( - resource_id=constants.RESOURCE_ID, - schema_title="system.Artifact", - uri=constants.MODEL_ARTIFACT_URI, - display_name=constants.DISPLAY_NAME, - schema_version=constants.SCHEMA_VERSION, - description=constants.DESCRIPTION, - metadata=constants.METADATA, - state=gca_artifact.Artifact.State.LIVE, - metadata_store_id="default", - project=constants.PROJECT, - location=constants.LOCATION, - credentials=None, - ) + mock_create_schema_base_artifact.assert_called_with( + project='abc', location='us-central1') assert artifact is mock_artifact diff --git a/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py b/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py index 54e9563a12..4be3022862 100644 --- a/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py +++ b/samples/model-builder/experiment_tracking/create_execution_with_sdk_sample_test.py @@ -14,13 +14,11 @@ import create_execution_with_sdk_sample -from google.cloud.aiplatform.compat.types import execution as gca_execution - import test_constants as constants def test_create_execution_sample( - mock_sdk_init, mock_create_artifact, mock_create_execution, mock_execution, + mock_sdk_init, mock_create_artifact, mock_create_schema_base_execution, mock_execution, ): input_art = mock_create_artifact() @@ -42,19 +40,7 @@ def test_create_execution_sample( project=constants.PROJECT, location=constants.LOCATION, ) - mock_create_execution.assert_called_with( - state=gca_execution.Execution.State.RUNNING, - schema_title="system.ContainerExecution", - resource_id=constants.RESOURCE_ID, - display_name=constants.DISPLAY_NAME, - schema_version=constants.SCHEMA_VERSION, - metadata=constants.METADATA, - description=constants.DESCRIPTION, - metadata_store_id="default", - project=None, - location=None, - credentials=None, - ) + mock_create_schema_base_execution.assert_called_with() mock_execution.assign_input_artifacts.assert_called_with([input_art]) mock_execution.assign_output_artifacts.assert_called_with([output_art]) diff --git a/tests/unit/aiplatform/test_metadata_schema.py b/tests/unit/aiplatform/test_metadata_schema.py index b133d5eebe..f550f61ab4 100644 --- a/tests/unit/aiplatform/test_metadata_schema.py +++ b/tests/unit/aiplatform/test_metadata_schema.py @@ -71,15 +71,61 @@ # artifact _TEST_ARTIFACT_ID = "test-artifact-id" -_TEST_ARTIFACT_NAME = f"{_TEST_PARENT}/artifacts/{_TEST_ARTIFACT_ID}" +_TEST_ARTIFACT_NAME = f"{_TEST_PARENT}/metadataStores/{_TEST_METADATA_STORE}/artifacts/{_TEST_ARTIFACT_ID}" # execution _TEST_EXECUTION_ID = "test-execution-id" -_TEST_EXECUTION_NAME = f"{_TEST_PARENT}/executions/{_TEST_EXECUTION_ID}" +_TEST_EXECUTION_NAME = f"{_TEST_PARENT}/metadataStores/{_TEST_METADATA_STORE}/executions/{_TEST_EXECUTION_ID}" # context _TEST_CONTEXT_ID = "test-context-id" -_TEST_CONTEXT_NAME = f"{_TEST_PARENT}/contexts/{_TEST_CONTEXT_ID}" +_TEST_CONTEXT_NAME = ( + f"{_TEST_PARENT}/metadataStores/{_TEST_METADATA_STORE}/contexts/{_TEST_CONTEXT_ID}" +) + + +@pytest.fixture +def get_artifact_mock(): + with patch.object(MetadataServiceClient, "get_artifact") as get_artifact_mock: + get_artifact_mock.return_value = GapicArtifact( + name=_TEST_ARTIFACT_NAME, + display_name=_TEST_DISPLAY_NAME, + schema_title=_TEST_SCHEMA_TITLE, + schema_version=_TEST_SCHEMA_VERSION, + description=_TEST_DESCRIPTION, + metadata=_TEST_METADATA, + state=GapicArtifact.State.STATE_UNSPECIFIED, + ) + yield get_artifact_mock + + +@pytest.fixture +def get_execution_mock(): + with patch.object(MetadataServiceClient, "get_execution") as get_execution_mock: + get_execution_mock.return_value = GapicExecution( + name=_TEST_EXECUTION_NAME, + display_name=_TEST_DISPLAY_NAME, + schema_title=_TEST_SCHEMA_TITLE, + schema_version=_TEST_SCHEMA_VERSION, + description=_TEST_DESCRIPTION, + metadata=_TEST_METADATA, + state=GapicExecution.State.RUNNING, + ) + yield get_execution_mock + + +@pytest.fixture +def get_context_mock(): + with patch.object(MetadataServiceClient, "get_context") as get_context_mock: + get_context_mock.return_value = GapicExecution( + name=_TEST_CONTEXT_NAME, + display_name=_TEST_DISPLAY_NAME, + schema_title=_TEST_SCHEMA_TITLE, + schema_version=_TEST_SCHEMA_VERSION, + description=_TEST_DESCRIPTION, + metadata=_TEST_METADATA, + ) + yield get_context_mock @pytest.fixture @@ -172,7 +218,7 @@ def test_base_class_without_schema_title_raises_error(self): with pytest.raises(TypeError): base_artifact.BaseArtifactSchema() - @pytest.mark.usefixtures("create_artifact_mock") + @pytest.mark.usefixtures("create_artifact_mock", "get_artifact_mock") def test_create_is_called_with_default_parameters(self, create_artifact_mock): aiplatform.init(project=_TEST_PROJECT) @@ -242,7 +288,7 @@ def test_base_class_without_schema_title_raises_error(self): with pytest.raises(TypeError): base_execution.BaseExecutionSchema() - @pytest.mark.usefixtures("create_execution_mock") + @pytest.mark.usefixtures("create_execution_mock", "get_execution_mock") def test_create_method_calls_gapic_library_with_correct_parameters( self, create_execution_mock ): @@ -394,7 +440,7 @@ def test_unmanaged_container_model_constructor_parameters_are_set_correctly(self metadata=_TEST_UPDATED_METADATA, ) expected_metadata = { - "test-param1": 2, + "test-param1": 2.0, "test-param2": "test-value-1", "test-param3": False, "predictSchemata": { @@ -409,7 +455,9 @@ def test_unmanaged_container_model_constructor_parameters_are_set_correctly(self assert artifact.uri == _TEST_URI assert artifact.display_name == _TEST_DISPLAY_NAME assert artifact.description == _TEST_DESCRIPTION - assert json.dumps(artifact.metadata) == json.dumps(expected_metadata) + assert json.dumps(artifact.metadata, sort_keys=True) == json.dumps( + expected_metadata, sort_keys=True + ) assert artifact.schema_version == _TEST_SCHEMA_VERSION @@ -489,7 +537,7 @@ def test_system_metrics_schema_title_is_set_correctly(self): def test_system_metrics_values_default_to_none(self): artifact = system_artifact_schema.Metrics() - assert artifact.metadata == {} + assert artifact._gca_resource.metadata is None def test_system_metrics_constructor_parameters_are_set_correctly(self): artifact = system_artifact_schema.Metrics( @@ -554,7 +602,7 @@ def teardown_method(self): initializer.global_pool.shutdown(wait=True) # Test system.Context Schemas - @pytest.mark.usefixtures("create_context_mock") + @pytest.mark.usefixtures("create_context_mock", "get_context_mock") def test_create_is_called_with_default_parameters(self, create_context_mock): aiplatform.init(project=_TEST_PROJECT) @@ -650,7 +698,7 @@ def test_container_spec_to_dict_method_returns_correct_schema(self): assert json.dumps(container_spec.to_dict()) == json.dumps(expected_results) - @pytest.mark.usefixtures("create_execution_mock") + @pytest.mark.usefixtures("create_execution_mock", "get_execution_mock") def test_start_execution_method_calls_gapic_library_with_correct_parameters( self, create_execution_mock ):