def get_metadata_store_mock():
with patch.object(
MetadataServiceClient, "get_metadata_store"
) as get_metadata_store_mock:
get_metadata_store_mock.return_value = GapicMetadataStore(
name=_TEST_METADATASTORE,
)
yield get_metadata_store_mock
def get_metadata_store_without_name_mock():
with patch.object(
MetadataServiceClient, "get_metadata_store"
) as get_metadata_store_mock:
get_metadata_store_mock.return_value = GapicMetadataStore(
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
yield get_metadata_store_mock
def get_default_metadata_store_mock():
with patch.object(MetadataServiceClient,
"get_metadata_store") as get_metadata_store_mock:
get_metadata_store_mock.return_value = GapicMetadataStore(
name=_TEST_DEFAULT_NAME,
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
yield get_metadata_store_mock
def create_metadata_store_mock():
with patch.object(
MetadataServiceClient, "create_metadata_store"
) as create_metadata_store_mock:
create_metadata_store_lro_mock = mock.Mock(operation.Operation)
create_metadata_store_lro_mock.result.return_value = GapicMetadataStore(
name=_TEST_METADATASTORE,
)
create_metadata_store_mock.return_value = create_metadata_store_lro_mock
yield create_metadata_store_mock
def get_metadata_store_mock_raise_not_found_exception():
with patch.object(
MetadataServiceClient, "get_metadata_store"
) as get_metadata_store_mock:
get_metadata_store_mock.side_effect = [
exceptions.NotFound("Test store not found."),
GapicMetadataStore(name=_TEST_METADATASTORE,),
]
yield get_metadata_store_mock
def create_metadata_store_mock():
with patch.object(MetadataServiceClient,
"create_metadata_store") as create_metadata_store_mock:
create_metadata_store_lro_mock = mock.Mock(operation.Operation)
create_metadata_store_lro_mock.result.return_value = GapicMetadataStore(
name=_TEST_NAME,
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
create_metadata_store_mock.return_value = create_metadata_store_lro_mock
yield create_metadata_store_mock
def test_create_non_default_metadata_store(self, create_metadata_store_mock):
aiplatform.init(project=_TEST_PROJECT)
my_metadata_store = metadata_store._MetadataStore._create(
metadata_store_id=_TEST_ID,
encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
)
expected_metadata_store = GapicMetadataStore(
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
create_metadata_store_mock.assert_called_once_with(
parent=_TEST_PARENT,
metadata_store_id=_TEST_ID,
metadata_store=expected_metadata_store,
)
expected_metadata_store.name = _TEST_NAME
assert my_metadata_store._gca_resource == expected_metadata_store