Repository: googleapis/java-datastore
Pull request #1409 (2024-04-25 14:38:48)

feat: have DatastoreException extend BaseHttpServiceException, convert gRPC status codes to existing Datastore codes
Building on top of changes in https://github.com/googleapis/java-datastore/pull/1400, this PR reverts the updates in https://github.com/googleapis/java-datastore/pull/1244 and instead translates gRPC error codes into existing Datastore ones.

This follows the same logic and patterns as in java-storage: https://github.com/googleapis/java-storage/blob/main/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcToHttpStatusCodeTranslation.java


Patch:
diff --git a/README.md b/README.md
index dcb40a9c..577bd7e0 100644
--- a/README.md
+++ b/README.md
@@ -57,13 +57,13 @@ implementation 'com.google.cloud:google-cloud-datastore'
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-datastore:2.19.0'
+implementation 'com.google.cloud:google-cloud-datastore:2.19.1'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.19.0"
+libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.19.1"
 ```
 <!-- {x-version-update-end} -->
 
@@ -380,7 +380,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
 [kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/java-datastore/java11.html
 [stability-image]: https://img.shields.io/badge/stability-stable-green
 [maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-datastore.svg
-[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-datastore/2.19.0
+[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-datastore/2.19.1
 [authentication]: https://github.com/googleapis/google-cloud-java#authentication
 [auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
 [predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java
index 5d01c3b9..d258feac 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java
@@ -16,16 +16,17 @@
 
 package com.google.cloud.datastore;
 
-import static com.google.cloud.BaseServiceException.isRetryable;
-
+import com.google.api.gax.grpc.GrpcStatusCode;
 import com.google.api.gax.rpc.ApiException;
-import com.google.api.gax.rpc.ErrorDetails;
+import com.google.api.gax.rpc.StatusCode;
 import com.google.cloud.BaseServiceException;
 import com.google.cloud.RetryHelper.RetryHelperException;
-import com.google.cloud.grpc.BaseGrpcServiceException;
+import com.google.cloud.http.BaseHttpServiceException;
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableSet;
+import io.grpc.StatusException;
+import io.grpc.StatusRuntimeException;
 import java.io.IOException;
-import java.util.Map;
 import java.util.Set;
 
 /**
@@ -34,7 +35,7 @@ import java.util.Set;
  * @see <a href="https://cloud.google.com/datastore/docs/concepts/errors#Error_Codes">Google Cloud
  *     Datastore error codes</a>
  */
-public final class DatastoreException extends BaseGrpcServiceException {
+public final class DatastoreException extends BaseHttpServiceException {
 
   // see https://cloud.google.com/datastore/docs/concepts/errors#Error_Codes"
   private static final Set<Error> RETRYABLE_ERRORS =
@@ -43,110 +44,106 @@ public final class DatastoreException extends BaseGrpcServiceException {
           new Error(4, "DEADLINE_EXCEEDED", false),
           new Error(14, "UNAVAILABLE", true));
   private static final long serialVersionUID = 2663750991205874435L;
-  private String reason;
-  private ApiException apiException;
 
   public DatastoreException(int code, String message, String reason) {
     this(code, message, reason, true, null);
-    this.reason = reason;
+  }
+
+  public DatastoreException(int code, String message, Throwable cause) {
+    super(code, message, null, true, RETRYABLE_ERRORS, cause);
   }
 
   public DatastoreException(int code, String message, String reason, Throwable cause) {
-    super(message, cause, code, isRetryable(code, reason, true, RETRYABLE_ERRORS));
-    this.reason = reason;
+    super(code, message, reason, true, RETRYABLE_ERRORS, cause);
   }
 
   public DatastoreException(
       int code, String message, String reason, boolean idempotent, Throwable cause) {
-    super(message, cause, code, isRetryable(code, reason, idempotent, RETRYABLE_ERRORS));
-    this.reason = reason;
+    super(code, message, reason, idempotent, RETRYABLE_ERRORS, cause);
   }
 
   public DatastoreException(IOException exception) {
-    super(exception, true);
-  }
-
-  public DatastoreException(ApiException apiException) {
-    super(apiException);
-    this.apiException = apiException;
+    super(exception, true, RETRYABLE_ERRORS);
   }
 
   /**
-   * Checks the underlying reason of the exception and if it's {@link ApiException} then return the
-   * specific domain otherwise null.
+   * Translate RetryHelperException to the DatastoreException that caused the error. This method
+   * will always throw an exception.
    *
-   * @see <a
-   *     href="https://github.com/googleapis/googleapis/blob/master/google/rpc/error_details.proto#L125">Domain</a>
-   * @return the logical grouping to which the "reason" belongs.
+   * @throws DatastoreException when {@code ex} was caused by a {@code DatastoreException}
    */
-  public String getDomain() {
-    if (this.apiException != null) {
-      return this.apiException.getDomain();
-    }
-    return null;
+  static DatastoreException translateAndThrow(RetryHelperException ex) {
+    BaseServiceException.translate(ex);
+    throw transformThrowable(ex);
   }
 
-  /**
-   * Checks the underlying reason of the exception and if it's {@link ApiException} then return a
-   * map of key-value pairs otherwise null.
-   *
-   * @see <a
-   *     href="https://github.com/googleapis/googleapis/blob/master/google/rpc/error_details.proto#L135">Metadata</a>
-   * @return the map of additional structured details about an error.
-   */
-  public Map<String, String> getMetadata() {
-    if (this.apiException != null) {
-      return this.apiException.getMetadata();
+  static BaseServiceException transformThrowable(Throwable t) {
+    if (t instanceof BaseServiceException) {
+      return (BaseServiceException) t;
     }
-    return null;
+    if (t.getCause() instanceof BaseServiceException) {
+      return (BaseServiceException) t.getCause();
+    }
+    if (t instanceof ApiException) {
+      return asDatastoreException((ApiException) t);
+    }
+    if (t.getCause() instanceof ApiException) {
+      return asDatastoreException((ApiException) t.getCause());
+    }
+    return getDatastoreException(t);
   }
 
-  /**
-   * Checks the underlying reason of the exception and if it's {@link ApiException} then return the
-   * ErrorDetails otherwise null.
-   *
-   * @see <a
-   *     href="https://github.com/googleapis/googleapis/blob/master/google/rpc/status.proto">Status</a>
-   * @see <a
-   *     href="https://github.com/googleapis/googleapis/blob/master/google/rpc/error_details.proto">Error
-   *     Details</a>
-   * @return An object containing getters for structured objects from error_details.proto.
-   */
-  public ErrorDetails getErrorDetails() {
-    if (this.apiException != null) {
-      return this.apiException.getErrorDetails();
+  private static DatastoreException getDatastoreException(Throwable t) {
+    // unwrap a RetryHelperException if that is what is being translated
+    if (t instanceof RetryHelperException) {
+      return new DatastoreException(UNKNOWN_CODE, t.getMessage(), null, t.getCause());
     }
-    return null;
+    return new DatastoreException(UNKNOWN_CODE, t.getMessage(), t);
   }
 
-  /**
-   * Checks the underlying reason of the exception and if it's {@link ApiException} then return the
-   * reason otherwise null/custom reason.
-   *
-   * @see <a
-   *     href="https://github.com/googleapis/googleapis/blob/master/google/rpc/error_details.proto#L117">Reason</a>
-   * @return the reason of an error.
-   */
-  @Override
-  public String getReason() {
-    if (this.apiException != null) {
-      return this.apiException.getReason();
+  static DatastoreException asDatastoreException(ApiException apiEx) {
+    int datastoreStatusCode = 0;
+    StatusCode statusCode = apiEx.getStatusCode();
+    if (statusCode instanceof GrpcStatusCode) {
+      GrpcStatusCode gsc = (GrpcStatusCode) statusCode;
+      datastoreStatusCode =
+          GrpcToDatastoreCodeTranslation.grpcCodeToDatastoreStatusCode(gsc.getTransportCode());
+    }
+
+    // If there is a gRPC exception in our cause, pull its error message up to be our
+    // message otherwise, create a generic error message with the status code.
+    String statusCodeName = statusCode.getCode().name();
+    String statusExceptionMessage = getStatusExceptionMessage(apiEx);
+
+    String message;
+    if (statusExceptionMessage != null) {
+      message = statusCodeName + ": " + statusExceptionMessage;
+    } else {
+      message = "Error: " + statusCodeName;
+    }
+
+    String reason = "";
+    if (Strings.isNullOrEmpty(apiEx.getReason())) {
+      if (apiEx.getStatusCode() != null) {
+        reason = apiEx.getStatusCode().getCode().name();
+      }
     }
-    return this.reason;
+    // It'd be better to use ExceptionData and BaseServiceException#<init>(ExceptionData) but,
+    // BaseHttpServiceException does not pass that through so we're stuck using this for now.
+    // TODO: When we can break the coupling to BaseHttpServiceException replace this
+    return new DatastoreException(datastoreStatusCode, message, reason, apiEx);
   }
 
-  /**
-   * Translate RetryHelperException to the DatastoreException that caused the error. This method
-   * will always throw an exception.
-   *
-   * @throws DatastoreException when {@code ex} was caused by a {@code DatastoreException}
-   */
-  static DatastoreException translateAndThrow(RetryHelperException ex) {
-    BaseServiceException.translate(ex);
-    if (ex.getCause() instanceof ApiException) {
-      throw new DatastoreException((ApiException) ex.getCause());
+  private static String getStatusExceptionMessage(Exception apiEx) {
+    if (apiEx.getMessage() != null) {
+      return apiEx.getMessage();
+    } else {
+      Throwable cause = apiEx.getCause();
+      if (cause instanceof StatusRuntimeException || cause instanceof StatusException) {
+        return cause.getMessage();
+      }
+      return null;
     }
-    throw new DatastoreException(UNKNOWN_CODE, ex.getMessage(), null, ex.getCause());
   }
 
   /**
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslation.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslation.java
new file mode 100644
index 00000000..1d63fb19
--- /dev/null
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslation.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.datastore;
+
+import com.google.api.gax.grpc.GrpcStatusCode;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.rpc.Code;
+import io.grpc.Status;
+import java.util.Map;
+import java.util.function.Function;
+
+final class GrpcToDatastoreCodeTranslation {
+  /** Mappings between gRPC status codes and their corresponding code numbers. */
+  private static final ImmutableList<StatusCodeMapping> STATUS_CODE_MAPPINGS =
+      ImmutableList.of(
+          StatusCodeMapping.of(Code.OK.getNumber(), Status.Code.OK),
+          StatusCodeMapping.of(Code.DATA_LOSS.getNumber(), Status.Code.DATA_LOSS),
+          StatusCodeMapping.of(Code.INVALID_ARGUMENT.getNumber(), Status.Code.INVALID_ARGUMENT),
+          StatusCodeMapping.of(Code.OUT_OF_RANGE.getNumber(), Status.Code.OUT_OF_RANGE),
+          StatusCodeMapping.of(Code.UNAUTHENTICATED.getNumber(), Status.Code.UNAUTHENTICATED),
+          StatusCodeMapping.of(Code.PERMISSION_DENIED.getNumber(), Status.Code.PERMISSION_DENIED),
+          StatusCodeMapping.of(Code.NOT_FOUND.getNumber(), Status.Code.NOT_FOUND),
+          StatusCodeMapping.of(Code.ALREADY_EXISTS.getNumber(), Status.Code.ALREADY_EXISTS),
+          StatusCodeMapping.of(
+              Code.FAILED_PRECONDITION.getNumber(), Status.Code.FAILED_PRECONDITION),
+          StatusCodeMapping.of(Code.RESOURCE_EXHAUSTED.getNumber(), Status.Code.RESOURCE_EXHAUSTED),
+          StatusCodeMapping.of(Code.INTERNAL.getNumber(), Status.Code.INTERNAL),
+          StatusCodeMapping.of(Code.UNIMPLEMENTED.getNumber(), Status.Code.UNIMPLEMENTED),
+          StatusCodeMapping.of(Code.UNAVAILABLE.getNumber(), Status.Code.UNAVAILABLE),
+          StatusCodeMapping.of(Code.DEADLINE_EXCEEDED.getNumber(), Status.Code.DEADLINE_EXCEEDED),
+          StatusCodeMapping.of(Code.ABORTED.getNumber(), Status.Code.ABORTED),
+          StatusCodeMapping.of(Code.CANCELLED.getNumber(), Status.Code.CANCELLED),
+          StatusCodeMapping.of(Code.UNKNOWN.getNumber(), Status.Code.UNKNOWN));
+
+  /** Index our {@link StatusCodeMapping} for constant time lookup by {@link Status.Code} */
+  private static final Map<Status.Code, StatusCodeMapping> GRPC_CODE_INDEX =
+      STATUS_CODE_MAPPINGS.stream()
+          .collect(
+              ImmutableMap.toImmutableMap(StatusCodeMapping::getGrpcCode, Function.identity()));
+
+  static int grpcCodeToDatastoreStatusCode(Status.Code code) {
+    StatusCodeMapping found = GRPC_CODE_INDEX.get(code);
+    // theoretically it's possible for gRPC to add a new code we haven't mapped here, if this
+    // happens fall through to our default of 0
+    if (found != null) {
+      return found.getDatastoreCode();
+    } else {
+      return 0;
+    }
+  }
+
+  /**
+   * Simple tuple class to bind together our corresponding http status code and {@link Status.Code}
+   * while providing easy access to the correct {@link GrpcStatusCode} where necessary.
+   */
+  private static final class StatusCodeMapping {
+
+    private final int datastoreCode;
+
+    private final Status.Code grpcCode;
+
+    private StatusCodeMapping(int datastoreCode, Status.Code grpcCode) {
+      this.datastoreCode = datastoreCode;
+      this.grpcCode = grpcCode;
+    }
+
+    public int getDatastoreCode() {
+      return datastoreCode;
+    }
+
+    public Status.Code getGrpcCode() {
+      return grpcCode;
+    }
+
+    static StatusCodeMapping of(int datastoreCode, Status.Code grpcCode) {
+      return new StatusCodeMapping(datastoreCode, grpcCode);
+    }
+  }
+}


Test patch:
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExceptionTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExceptionTest.java
index 33f5ebb9..8c52b551 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExceptionTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreExceptionTest.java
@@ -16,9 +16,6 @@
 
 package com.google.cloud.datastore;
 
-import static com.google.common.truth.Truth.assertThat;
-import static com.google.rpc.Code.FAILED_PRECONDITION;
-import static java.util.Collections.singletonMap;
 import static org.easymock.EasyMock.createMock;
 import static org.easymock.EasyMock.expect;
 import static org.easymock.EasyMock.replay;
@@ -27,20 +24,11 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
-import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import com.google.api.gax.grpc.GrpcStatusCode;
-import com.google.api.gax.rpc.ApiException;
-import com.google.api.gax.rpc.ApiExceptionFactory;
-import com.google.api.gax.rpc.ErrorDetails;
-import com.google.api.gax.rpc.StatusCode;
 import com.google.cloud.BaseServiceException;
 import com.google.cloud.RetryHelper;
-import com.google.protobuf.Any;
-import com.google.rpc.ErrorInfo;
-import io.grpc.Status;
 import java.io.IOException;
 import java.net.SocketTimeoutException;
 import org.junit.Test;
@@ -88,78 +76,39 @@ public class DatastoreExceptionTest {
     assertEquals("message", exception.getMessage());
     assertFalse(exception.isRetryable());
     assertSame(cause, exception.getCause());
-
-    exception = new DatastoreException(2, "message", "INTERNAL", true, cause);
-    assertEquals(2, exception.getCode());
-    assertEquals("INTERNAL", exception.getReason());
-    assertEquals("message", exception.getMessage());
-    assertFalse(exception.isRetryable());
-    assertSame(cause, exception.getCause());
-
-    ApiException apiException = createApiException();
-    exception = new DatastoreException(apiException);
-    assertEquals(400, exception.getCode());
-    assertEquals("MISSING_INDEXES", exception.getReason());
-    assertThat(exception.getMetadata())
-        .isEqualTo(singletonMap("missing_indexes_url", "__some__url__"));
-    assertSame(apiException, exception.getCause());
-  }
-
-  @Test
-  public void testApiException() {
-    ApiException apiException = createApiException();
-    DatastoreException datastoreException = new DatastoreException(apiException);
-
-    assertThat(datastoreException.getReason()).isEqualTo("MISSING_INDEXES");
-    assertThat(datastoreException.getDomain()).isEqualTo("datastore.googleapis.com");
-    assertThat(datastoreException.getMetadata())
-        .isEqualTo(singletonMap("missing_indexes_url", "__some__url__"));
-    assertThat(datastoreException.getErrorDetails()).isEqualTo(apiException.getErrorDetails());
   }
 
   @Test
   public void testTranslateAndThrow() {
     Exception cause = new DatastoreException(14, "message", "UNAVAILABLE");
-    final RetryHelper.RetryHelperException exceptionMock =
+    RetryHelper.RetryHelperException exceptionMock =
         createMock(RetryHelper.RetryHelperException.class);
     expect(exceptionMock.getCause()).andReturn(cause).times(2);
     replay(exceptionMock);
-    BaseServiceException ex =
-        assertThrows(
-            BaseServiceException.class, () -> DatastoreException.translateAndThrow(exceptionMock));
-    assertEquals(14, ex.getCode());
-    assertEquals("message", ex.getMessage());
-    assertTrue(ex.isRetryable());
-    verify(exceptionMock);
-
-    cause = createApiException();
-    final RetryHelper.RetryHelperException exceptionMock2 =
-        createMock(RetryHelper.RetryHelperException.class);
-    expect(exceptionMock2.getCause()).andReturn(cause).times(3);
-    replay(exceptionMock2);
-    DatastoreException ex2 =
-        assertThrows(
-            DatastoreException.class, () -> DatastoreException.translateAndThrow(exceptionMock2));
-    assertThat(ex2.getReason()).isEqualTo("MISSING_INDEXES");
-    assertThat(ex2.getDomain()).isEqualTo("datastore.googleapis.com");
-    assertThat(ex2.getMetadata()).isEqualTo(singletonMap("missing_indexes_url", "__some__url__"));
-    assertThat(ex2.getErrorDetails()).isEqualTo(((ApiException) cause).getErrorDetails());
-    verify(exceptionMock2);
-
+    try {
+      DatastoreException.translateAndThrow(exceptionMock);
+    } catch (BaseServiceException ex) {
+      assertEquals(14, ex.getCode());
+      assertEquals("message", ex.getMessage());
+      assertTrue(ex.isRetryable());
+    } finally {
+      verify(exceptionMock);
+    }
     cause = new IllegalArgumentException("message");
-    final RetryHelper.RetryHelperException exceptionMock3 =
-        createMock(RetryHelper.RetryHelperException.class);
-    expect(exceptionMock3.getMessage()).andReturn("message").times(1);
-    expect(exceptionMock3.getCause()).andReturn(cause).times(3);
-    replay(exceptionMock3);
-    BaseServiceException ex3 =
-        assertThrows(
-            BaseServiceException.class, () -> DatastoreException.translateAndThrow(exceptionMock3));
-    assertEquals(DatastoreException.UNKNOWN_CODE, ex3.getCode());
-    assertEquals("message", ex3.getMessage());
-    assertFalse(ex3.isRetryable());
-    assertSame(cause, ex3.getCause());
-    verify(exceptionMock3);
+    exceptionMock = createMock(RetryHelper.RetryHelperException.class);
+    expect(exceptionMock.getMessage()).andReturn("message").times(1);
+    expect(exceptionMock.getCause()).andReturn(cause).times(4);
+    replay(exceptionMock);
+    try {
+      DatastoreException.translateAndThrow(exceptionMock);
+    } catch (BaseServiceException ex) {
+      assertEquals(DatastoreException.UNKNOWN_CODE, ex.getCode());
+      assertEquals("message", ex.getMessage());
+      assertFalse(ex.isRetryable());
+      assertSame(cause, ex.getCause());
+    } finally {
+      verify(exceptionMock);
+    }
   }
 
   @Test
@@ -172,26 +121,4 @@ public class DatastoreExceptionTest {
       assertEquals("message a 1", ex.getMessage());
     }
   }
-
-  private ApiException createApiException() {
-    // Simulating google.rpc.Status with an ErrorInfo
-    ErrorInfo errorInfo =
-        ErrorInfo.newBuilder()
-            .setDomain("datastore.googleapis.com")
-            .setReason("MISSING_INDEXES")
-            .putMetadata("missing_indexes_url", "__some__url__")
-            .build();
-    com.google.rpc.Status status =
-        com.google.rpc.Status.newBuilder()
-            .setCode(FAILED_PRECONDITION.getNumber())
-            .setMessage("The query requires indexes.")
-            .addDetails(Any.pack(errorInfo))
-            .build();
-
-    // Using Gax to convert to ApiException
-    StatusCode statusCode = GrpcStatusCode.of(Status.fromCodeValue(status.getCode()).getCode());
-    ErrorDetails errorDetails =
-        ErrorDetails.builder().setRawErrorMessages(status.getDetailsList()).build();
-    return ApiExceptionFactory.createException(null, statusCode, true, errorDetails);
-  }
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslationTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslationTest.java
new file mode 100644
index 00000000..3f297989
--- /dev/null
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/GrpcToDatastoreCodeTranslationTest.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.datastore;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import io.grpc.Status.Code;
+import java.util.EnumMap;
+import org.junit.Test;
+
+public class GrpcToDatastoreCodeTranslationTest {
+  @Test
+  public void grpcCodeToDatastoreCode_expectedMapping() {
+    EnumMap<Code, Integer> expected = new EnumMap<>(Code.class);
+    expected.put(Code.OK, com.google.rpc.Code.OK.getNumber());
+    expected.put(Code.INVALID_ARGUMENT, com.google.rpc.Code.INVALID_ARGUMENT.getNumber());
+    expected.put(Code.OUT_OF_RANGE, com.google.rpc.Code.OUT_OF_RANGE.getNumber());
+    expected.put(Code.UNAUTHENTICATED, com.google.rpc.Code.UNAUTHENTICATED.getNumber());
+    expected.put(Code.PERMISSION_DENIED, com.google.rpc.Code.PERMISSION_DENIED.getNumber());
+    expected.put(Code.NOT_FOUND, com.google.rpc.Code.NOT_FOUND.getNumber());
+    expected.put(Code.FAILED_PRECONDITION, com.google.rpc.Code.FAILED_PRECONDITION.getNumber());
+    expected.put(Code.ALREADY_EXISTS, com.google.rpc.Code.ALREADY_EXISTS.getNumber());
+    expected.put(Code.RESOURCE_EXHAUSTED, com.google.rpc.Code.RESOURCE_EXHAUSTED.getNumber());
+    expected.put(Code.INTERNAL, com.google.rpc.Code.INTERNAL.getNumber());
+    expected.put(Code.UNIMPLEMENTED, com.google.rpc.Code.UNIMPLEMENTED.getNumber());
+    expected.put(Code.UNAVAILABLE, com.google.rpc.Code.UNAVAILABLE.getNumber());
+    expected.put(Code.ABORTED, com.google.rpc.Code.ABORTED.getNumber());
+    expected.put(Code.CANCELLED, com.google.rpc.Code.CANCELLED.getNumber());
+    expected.put(Code.UNKNOWN, com.google.rpc.Code.UNKNOWN.getNumber());
+    expected.put(Code.DEADLINE_EXCEEDED, com.google.rpc.Code.DEADLINE_EXCEEDED.getNumber());
+    expected.put(Code.DATA_LOSS, com.google.rpc.Code.DATA_LOSS.getNumber());
+
+    EnumMap<Code, Integer> actual = new EnumMap<>(Code.class);
+    for (Code c : Code.values()) {
+      actual.put(c, GrpcToDatastoreCodeTranslation.grpcCodeToDatastoreStatusCode(c));
+    }
+
+    assertThat(actual).isEqualTo(expected);
+  }
+}
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITDatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITDatastoreTest.java
index 7c105672..ac9587b2 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITDatastoreTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITDatastoreTest.java
@@ -16,7 +16,11 @@
 
 package com.google.cloud.datastore.it;
 
+import static com.google.api.gax.rpc.StatusCode.Code.ALREADY_EXISTS;
+import static com.google.api.gax.rpc.StatusCode.Code.DEADLINE_EXCEEDED;
+import static com.google.api.gax.rpc.StatusCode.Code.FAILED_PRECONDITION;
 import static com.google.api.gax.rpc.StatusCode.Code.INVALID_ARGUMENT;
+import static com.google.api.gax.rpc.StatusCode.Code.NOT_FOUND;
 import static com.google.cloud.datastore.aggregation.Aggregation.avg;
 import static com.google.cloud.datastore.aggregation.Aggregation.count;
 import static com.google.cloud.datastore.aggregation.Aggregation.sum;
@@ -28,8 +32,8 @@ import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 
 import com.google.cloud.Timestamp;
 import com.google.cloud.Tuple;
@@ -70,9 +74,9 @@ import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
 import com.google.cloud.datastore.TimestampValue;
 import com.google.cloud.datastore.Transaction;
 import com.google.cloud.datastore.ValueType;
-import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.truth.Truth;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.datastore.v1.TransactionOptions.ReadOnly;
 import java.util.ArrayList;
@@ -383,19 +387,8 @@ public abstract class AbstractITDatastoreTest {
     assertEquals(ENTITY3, list.get(2));
     assertEquals(3, list.size());
 
-    try {
-      transaction.commit();
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      assertEquals("FAILED_PRECONDITION", expected.getReason());
-    }
-
-    try {
-      transaction.rollback();
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      assertEquals("FAILED_PRECONDITION", expected.getReason());
-    }
+    DatastoreException expected = assertThrows(DatastoreException.class, transaction::commit);
+    assertDatastoreException(expected, FAILED_PRECONDITION.name(), 0);
   }
 
   @Test
@@ -492,12 +485,8 @@ public abstract class AbstractITDatastoreTest {
     transaction.rollback();
     transaction.rollback(); // should be safe to repeat rollback calls
 
-    try {
-      transaction.commit();
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      assertEquals("FAILED_PRECONDITION", expected.getReason());
-    }
+    DatastoreException expected = assertThrows(DatastoreException.class, transaction::commit);
+    assertDatastoreException(expected, FAILED_PRECONDITION.name(), 0);
 
     List<Entity> list = datastore.fetch(KEY1, KEY2, KEY3);
     assertEquals(ENTITY1, list.get(0));
@@ -544,12 +533,8 @@ public abstract class AbstractITDatastoreTest {
     assertEquals(PARTIAL_ENTITY3.getNames(), datastore.get(generatedKeys.get(0)).getNames());
     assertEquals(PARTIAL_ENTITY3.getKey(), IncompleteKey.newBuilder(generatedKeys.get(0)).build());
 
-    try {
-      batch.submit();
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      assertEquals("FAILED_PRECONDITION", expected.getReason());
-    }
+    DatastoreException expected = assertThrows(DatastoreException.class, batch::submit);
+    assertDatastoreException(expected, FAILED_PRECONDITION.name(), 0);
 
     batch = datastore.newBatch();
     batch.delete(entity4.getKey(), entity5.getKey(), entity6.getKey());
@@ -1256,12 +1241,11 @@ public abstract class AbstractITDatastoreTest {
     assertEquals(EMPTY_LIST_VALUE, entity3.getValue("emptyList"));
     assertEquals(8, entity3.getNames().size());
     assertFalse(entity3.contains("bla"));
-    try {
-      entity3.getString("str");
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      // expected - no such property
-    }
+
+    DatastoreException expected =
+        assertThrows(DatastoreException.class, () -> entity3.getString("str"));
+    assertDatastoreException(expected, FAILED_PRECONDITION.name(), 0);
+
     assertFalse(result.hasNext());
     datastore.delete(ENTITY3.getKey());
   }
@@ -1273,12 +1257,9 @@ public abstract class AbstractITDatastoreTest {
     assertNull(keys.get(1));
     assertEquals(2, keys.size());
 
-    try {
-      datastore.add(ENTITY1);
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      // expected;
-    }
+    DatastoreException expected =
+        assertThrows(DatastoreException.class, () -> datastore.add(ENTITY1));
+    assertDatastoreException(expected, ALREADY_EXISTS.name(), 6);
 
     List<Entity> entities = datastore.add(ENTITY3, PARTIAL_ENTITY1, PARTIAL_ENTITY2);
     assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
@@ -1301,12 +1282,10 @@ public abstract class AbstractITDatastoreTest {
     assertNull(keys.get(1));
     assertEquals(2, keys.size());
 
-    try {
-      datastore.update(ENTITY3);
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      // expected;
-    }
+    DatastoreException expected =
+        assertThrows(DatastoreException.class, () -> datastore.update(ENTITY3));
+    assertDatastoreException(expected, NOT_FOUND.name(), 5);
+
     datastore.add(ENTITY3);
     assertEquals(ENTITY3, datastore.get(ENTITY3.getKey()));
     Entity entity3 = Entity.newBuilder(ENTITY3).clear().set("bla", new NullValue()).build();
@@ -1316,6 +1295,12 @@ public abstract class AbstractITDatastoreTest {
     datastore.delete(ENTITY3.getKey());
   }
 
+  private void assertDatastoreException(
+      DatastoreException expected, String reason, int datastoreStatusCode) {
+    Truth.assertThat(expected.getReason()).isEqualTo(reason);
+    Truth.assertThat(expected.getCode()).isEqualTo(datastoreStatusCode);
+  }
+
   @Test
   public void testPut() {
     Entity updatedEntity = Entity.newBuilder(ENTITY1).set("new_property", 42L).build();
@@ -1392,12 +1377,9 @@ public abstract class AbstractITDatastoreTest {
           }
         };
 
-    try {
-      datastore.runInTransaction(callable2);
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      assertEquals(4, ((DatastoreException) expected.getCause()).getCode());
-    }
+    DatastoreException expected =
+        assertThrows(DatastoreException.class, () -> datastore.runInTransaction(callable2));
+    assertDatastoreException((DatastoreException) expected.getCause(), DEADLINE_EXCEEDED.name(), 4);
   }
 
   @Test
@@ -1446,18 +1428,10 @@ public abstract class AbstractITDatastoreTest {
             .setReadOnly(TransactionOptions.ReadOnly.getDefaultInstance())
             .build();
 
-    try {
-      datastore.runInTransaction(callable2, readOnlyOptions);
-      fail("Expecting a failure");
-    } catch (DatastoreException expected) {
-      if (datastore.getOptions().getTransportOptions() instanceof GrpcTransportOptions) {
-        assertEquals(
-            INVALID_ARGUMENT.getHttpStatusCode(),
-            ((DatastoreException) expected.getCause()).getCode());
-      } else {
-        assertEquals(3, ((DatastoreException) expected.getCause()).getCode());
-      }
-    }
+    DatastoreException expected =
+        assertThrows(
+            DatastoreException.class, () -> datastore.runInTransaction(callable2, readOnlyOptions));
+    assertDatastoreException((DatastoreException) expected.getCause(), INVALID_ARGUMENT.name(), 3);
   }
 
   @Test