def test_update_execution(self, update_execution_mock):
aiplatform.init(project=_TEST_PROJECT)
my_execution = execution._Execution._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,
)
my_execution.update(_TEST_UPDATED_METADATA)
updated_execution = GapicExecution(
name=_TEST_EXECUTION_NAME,
schema_title=_TEST_SCHEMA_TITLE,
schema_version=_TEST_SCHEMA_VERSION,
display_name=_TEST_DISPLAY_NAME,
description=_TEST_DESCRIPTION,
metadata=_TEST_UPDATED_METADATA,
)
update_execution_mock.assert_called_once_with(
execution=updated_execution)
assert my_execution._gca_resource == updated_execution
def test_get_or_create_execution(self,
get_execution_for_get_or_create_mock,
create_execution_mock):
aiplatform.init(project=_TEST_PROJECT)
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,
)
expected_execution = GapicExecution(
schema_title=_TEST_SCHEMA_TITLE,
schema_version=_TEST_SCHEMA_VERSION,
display_name=_TEST_DISPLAY_NAME,
description=_TEST_DESCRIPTION,
metadata=_TEST_METADATA,
)
get_execution_for_get_or_create_mock.assert_called_once_with(
name=_TEST_EXECUTION_NAME)
create_execution_mock.assert_called_once_with(
parent=_TEST_PARENT,
execution_id=_TEST_EXECUTION_ID,
execution=expected_execution,
)
expected_execution.name = _TEST_EXECUTION_NAME
assert my_execution._gca_resource == expected_execution
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_RUN,
schema_title=constants.SYSTEM_RUN,
schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN],
)
yield get_execution_mock
def list_executions_mock():
with patch.object(MetadataServiceClient, "list_executions") as list_executions_mock:
list_executions_mock.return_value = [
GapicExecution(
name=_TEST_EXECUTION_NAME,
display_name=_TEST_RUN,
schema_title=constants.SYSTEM_RUN,
schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN],
metadata=_TEST_PARAMS,
),
GapicExecution(
name=_TEST_OTHER_EXECUTION_NAME,
display_name=_TEST_OTHER_RUN,
schema_title=constants.SYSTEM_RUN,
schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN],
metadata=_TEST_OTHER_PARAMS,
),
]
yield list_executions_mock
def update_execution_mock():
with patch.object(MetadataServiceClient,
"update_execution") as update_execution_mock:
update_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_UPDATED_METADATA,
)
yield update_execution_mock
def list_executions_mock():
with patch.object(MetadataServiceClient,
"list_executions") as list_executions_mock:
list_executions_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,
),
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,
),
]
yield list_executions_mock
def test_log_params(
self, update_execution_mock,
):
aiplatform.init(
project=_TEST_PROJECT, location=_TEST_LOCATION, experiment=_TEST_EXPERIMENT
)
aiplatform.start_run(_TEST_RUN)
aiplatform.log_params(_TEST_PARAMS)
updated_execution = GapicExecution(
name=_TEST_EXECUTION_NAME,
display_name=_TEST_RUN,
schema_title=constants.SYSTEM_RUN,
schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_RUN],
metadata=_TEST_PARAMS,
)
update_execution_mock.assert_called_once_with(execution=updated_execution)
def test_list_executions(self, list_executions_mock):
aiplatform.init(project=_TEST_PROJECT)
filter = "test-filter"
execution_list = execution._Execution.list(
filter=filter, metadata_store_id=_TEST_METADATA_STORE)
expected_execution = GapicExecution(
name=_TEST_EXECUTION_NAME,
schema_title=_TEST_SCHEMA_TITLE,
schema_version=_TEST_SCHEMA_VERSION,
display_name=_TEST_DISPLAY_NAME,
description=_TEST_DESCRIPTION,
metadata=_TEST_METADATA,
)
list_executions_mock.assert_called_once_with(
request=ListExecutionsRequest(
parent=_TEST_PARENT,
filter=filter,
))
assert len(execution_list) == 2
assert execution_list[0]._gca_resource == expected_execution
assert execution_list[1]._gca_resource == expected_execution