def get_tensorboard_mock():
with patch.object(tensorboard_service_client.TensorboardServiceClient,
"get_tensorboard") as get_tensorboard_mock:
get_tensorboard_mock.return_value = gca_tensorboard.Tensorboard(
name=_TEST_NAME,
display_name=_TEST_DISPLAY_NAME,
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
yield get_tensorboard_mock
def create_tensorboard_mock():
with patch.object(tensorboard_service_client.TensorboardServiceClient,
"create_tensorboard") as create_tensorboard_mock:
create_tensorboard_lro_mock = mock.Mock(operation.Operation)
create_tensorboard_lro_mock.result.return_value = gca_tensorboard.Tensorboard(
name=_TEST_NAME,
display_name=_TEST_DISPLAY_NAME,
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
create_tensorboard_mock.return_value = create_tensorboard_lro_mock
yield create_tensorboard_mock
def test_update_tensorboard_encryption_spec(self, update_tensorboard_mock):
aiplatform.init(project=_TEST_PROJECT)
my_tensorboard = tensorboard.Tensorboard(tensorboard_name=_TEST_NAME)
my_tensorboard.update(encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME)
expected_tensorboard = gca_tensorboard.Tensorboard(
name=_TEST_NAME,
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
update_tensorboard_mock.assert_called_once_with(
update_mask=field_mask_pb2.FieldMask(paths=["encryption_spec"]),
tensorboard=expected_tensorboard,
metadata=_TEST_REQUEST_METADATA,
)
def test_update_tensorboard_display_name(self, update_tensorboard_mock):
aiplatform.init(project=_TEST_PROJECT)
my_tensorboard = tensorboard.Tensorboard(tensorboard_name=_TEST_NAME)
my_tensorboard.update(display_name=_TEST_DISPLAY_NAME_UPDATE)
expected_tensorboard = gca_tensorboard.Tensorboard(
name=_TEST_NAME,
display_name=_TEST_DISPLAY_NAME_UPDATE,
)
update_tensorboard_mock.assert_called_once_with(
update_mask=field_mask_pb2.FieldMask(paths=["display_name"]),
tensorboard=expected_tensorboard,
metadata=_TEST_REQUEST_METADATA,
)
def test_create_tensorboard(self, create_tensorboard_mock):
aiplatform.init(project=_TEST_PROJECT, )
tensorboard.Tensorboard.create(
display_name=_TEST_DISPLAY_NAME,
encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
)
expected_tensorboard = gca_tensorboard.Tensorboard(
display_name=_TEST_DISPLAY_NAME,
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
create_tensorboard_mock.assert_called_once_with(
parent=_TEST_PARENT,
tensorboard=expected_tensorboard,
metadata=_TEST_REQUEST_METADATA,
)
