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, )