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
