def query_execution_inputs_and_outputs_mock(): with patch.object( MetadataServiceClient, "query_execution_inputs_and_outputs" ) as query_execution_inputs_and_outputs_mock: query_execution_inputs_and_outputs_mock.side_effect = [ LineageSubgraph( artifacts=[ GapicArtifact( name=_TEST_ARTIFACT_NAME, display_name=_TEST_ARTIFACT_ID, schema_title=constants.SYSTEM_METRICS, schema_version=constants.SCHEMA_VERSIONS[ constants.SYSTEM_METRICS ], metadata=_TEST_METRICS, ), ], ), LineageSubgraph( artifacts=[ GapicArtifact( name=_TEST_OTHER_ARTIFACT_NAME, display_name=_TEST_OTHER_ARTIFACT_ID, schema_title=constants.SYSTEM_METRICS, schema_version=constants.SCHEMA_VERSIONS[ constants.SYSTEM_METRICS ], metadata=_TEST_OTHER_METRICS, ), ], ), ] yield query_execution_inputs_and_outputs_mock
def test_update_artifact(self, update_artifact_mock): aiplatform.init(project=_TEST_PROJECT) my_artifact = artifact._Artifact._create( resource_id=_TEST_ARTIFACT_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, schema_version=_TEST_SCHEMA_VERSION, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, metadata_store_id=_TEST_METADATA_STORE, ) my_artifact.update(_TEST_UPDATED_METADATA) updated_artifact = GapicArtifact( name=_TEST_ARTIFACT_NAME, schema_title=_TEST_SCHEMA_TITLE, schema_version=_TEST_SCHEMA_VERSION, display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_UPDATED_METADATA, ) update_artifact_mock.assert_called_once_with(artifact=updated_artifact) assert my_artifact._gca_resource == updated_artifact
def test_get_or_create_artifact(self, get_artifact_for_get_or_create_mock, create_artifact_mock): aiplatform.init(project=_TEST_PROJECT) my_artifact = artifact._Artifact.get_or_create( resource_id=_TEST_ARTIFACT_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, schema_version=_TEST_SCHEMA_VERSION, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, metadata_store_id=_TEST_METADATA_STORE, ) expected_artifact = GapicArtifact( schema_title=_TEST_SCHEMA_TITLE, schema_version=_TEST_SCHEMA_VERSION, display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, ) get_artifact_for_get_or_create_mock.assert_called_once_with( name=_TEST_ARTIFACT_NAME) create_artifact_mock.assert_called_once_with( parent=_TEST_PARENT, artifact_id=_TEST_ARTIFACT_ID, artifact=expected_artifact, ) expected_artifact.name = _TEST_ARTIFACT_NAME assert my_artifact._gca_resource == expected_artifact
def test_query_input_and_output_artifacts( self, query_execution_inputs_and_outputs_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) my_execution = execution._Execution.get_or_create( resource_id=_TEST_EXECUTION_ID, schema_title=_TEST_SCHEMA_TITLE, display_name=_TEST_DISPLAY_NAME, schema_version=_TEST_SCHEMA_VERSION, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, metadata_store_id=_TEST_METADATA_STORE, ) artifact_list = my_execution.query_input_and_output_artifacts() expected_artifact = GapicArtifact( name=_TEST_ARTIFACT_NAME, schema_title=_TEST_SCHEMA_TITLE, schema_version=_TEST_SCHEMA_VERSION, display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, ) query_execution_inputs_and_outputs_mock.assert_called_once_with( execution=_TEST_EXECUTION_NAME, ) assert len(artifact_list) == 2 assert artifact_list[0]._gca_resource == expected_artifact assert artifact_list[1]._gca_resource == expected_artifact
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_ARTIFACT_ID, schema_title=constants.SYSTEM_METRICS, schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_METRICS], ) yield get_artifact_mock
def update_artifact_mock(): with patch.object(MetadataServiceClient, "update_artifact") as update_artifact_mock: update_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_UPDATED_METADATA, ) yield update_artifact_mock
def list_artifacts_mock(): with patch.object(MetadataServiceClient, "list_artifacts") as list_artifacts_mock: list_artifacts_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, ), 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, ), ] yield list_artifacts_mock
def query_execution_inputs_and_outputs_mock(): with patch.object(MetadataServiceClient, "query_execution_inputs_and_outputs" ) as query_execution_inputs_and_outputs_mock: query_execution_inputs_and_outputs_mock.return_value = LineageSubgraph( artifacts=[ 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, ), 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, ), ], ) yield query_execution_inputs_and_outputs_mock
def test_log_metrics( self, update_artifact_mock, ): aiplatform.init( project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT ) aiplatform.start_run(_TEST_RUN) aiplatform.log_metrics(_TEST_METRICS) updated_artifact = GapicArtifact( name=_TEST_ARTIFACT_NAME, display_name=_TEST_ARTIFACT_ID, schema_title=constants.SYSTEM_METRICS, schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_METRICS], metadata=_TEST_METRICS, ) update_artifact_mock.assert_called_once_with(artifact=updated_artifact)
def test_list_artifacts(self, list_artifacts_mock): aiplatform.init(project=_TEST_PROJECT) filter = "test-filter" artifact_list = artifact._Artifact.list( filter=filter, metadata_store_id=_TEST_METADATA_STORE) expected_artifact = GapicArtifact( name=_TEST_ARTIFACT_NAME, schema_title=_TEST_SCHEMA_TITLE, schema_version=_TEST_SCHEMA_VERSION, display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, metadata=_TEST_METADATA, ) list_artifacts_mock.assert_called_once_with( request=ListArtifactsRequest( parent=_TEST_PARENT, filter=filter, )) assert len(artifact_list) == 2 assert artifact_list[0]._gca_resource == expected_artifact assert artifact_list[1]._gca_resource == expected_artifact