def test_update_context(self, update_context_mock):
aiplatform.init(project=_TEST_PROJECT)
my_context = context._Context._create(
resource_id=_TEST_CONTEXT_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_context.update(_TEST_UPDATED_METADATA)
updated_context = GapicContext(
name=_TEST_CONTEXT_NAME,
schema_title=_TEST_SCHEMA_TITLE,
schema_version=_TEST_SCHEMA_VERSION,
display_name=_TEST_DISPLAY_NAME,
description=_TEST_DESCRIPTION,
metadata=_TEST_UPDATED_METADATA,
)
update_context_mock.assert_called_once_with(context=updated_context)
assert my_context._gca_resource == updated_context
def test_get_or_create_context(self, get_context_for_get_or_create_mock,
create_context_mock):
aiplatform.init(project=_TEST_PROJECT)
my_context = context._Context.get_or_create(
resource_id=_TEST_CONTEXT_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_context = GapicContext(
schema_title=_TEST_SCHEMA_TITLE,
schema_version=_TEST_SCHEMA_VERSION,
display_name=_TEST_DISPLAY_NAME,
description=_TEST_DESCRIPTION,
metadata=_TEST_METADATA,
)
get_context_for_get_or_create_mock.assert_called_once_with(
name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY)
create_context_mock.assert_called_once_with(
parent=_TEST_PARENT,
context_id=_TEST_CONTEXT_ID,
context=expected_context,
)
expected_context.name = _TEST_CONTEXT_NAME
assert my_context._gca_resource == expected_context
def update_context_mock():
with patch.object(MetadataServiceClient, "update_context") as update_context_mock:
update_context_mock.return_value = GapicContext(
name=_TEST_CONTEXT_NAME,
display_name=_TEST_EXPERIMENT,
description=_TEST_OTHER_EXPERIMENT_DESCRIPTION,
schema_title=constants.SYSTEM_EXPERIMENT,
schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT],
metadata=constants.EXPERIMENT_METADATA,
)
yield update_context_mock
def get_context_mock():
with patch.object(MetadataServiceClient,
"get_context") as get_context_mock:
get_context_mock.return_value = GapicContext(
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
def get_pipeline_context_mock():
with patch.object(
MetadataServiceClient, "get_context"
) as get_pipeline_context_mock:
get_pipeline_context_mock.return_value = GapicContext(
name=_TEST_CONTEXT_NAME,
display_name=_TEST_EXPERIMENT,
schema_title=constants.SYSTEM_PIPELINE,
schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_PIPELINE],
metadata=constants.EXPERIMENT_METADATA,
)
yield get_pipeline_context_mock
def test_init_experiment_without_existing_description(self, update_context_mock):
aiplatform.init(
project=_TEST_PROJECT,
location=_TEST_LOCATION,
experiment=_TEST_EXPERIMENT,
experiment_description=_TEST_OTHER_EXPERIMENT_DESCRIPTION,
)
experiment_context = GapicContext(
name=_TEST_CONTEXT_NAME,
display_name=_TEST_EXPERIMENT,
description=_TEST_OTHER_EXPERIMENT_DESCRIPTION,
schema_title=constants.SYSTEM_EXPERIMENT,
schema_version=constants.SCHEMA_VERSIONS[constants.SYSTEM_EXPERIMENT],
metadata=constants.EXPERIMENT_METADATA,
)
update_context_mock.assert_called_once_with(context=experiment_context)
def test_list_contexts(self, list_contexts_mock):
aiplatform.init(project=_TEST_PROJECT)
filter = "test-filter"
context_list = context._Context.list(
filter=filter, metadata_store_id=_TEST_METADATA_STORE)
expected_context = GapicContext(
name=_TEST_CONTEXT_NAME,
schema_title=_TEST_SCHEMA_TITLE,
schema_version=_TEST_SCHEMA_VERSION,
display_name=_TEST_DISPLAY_NAME,
description=_TEST_DESCRIPTION,
metadata=_TEST_METADATA,
)
list_contexts_mock.assert_called_once_with(request=ListContextsRequest(
parent=_TEST_PARENT,
filter=filter,
))
assert len(context_list) == 2
assert context_list[0]._gca_resource == expected_context
assert context_list[1]._gca_resource == expected_context