diff --git a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/BaseDatastoreBatchWriter.java b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/BaseDatastoreBatchWriter.java
index 8804ef7e8dda..59a2a033d015 100644
--- a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/BaseDatastoreBatchWriter.java
+++ b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/BaseDatastoreBatchWriter.java
@@ -131,17 +131,65 @@ public final void update(Entity... entities) {
     }
   }
 
-  @SafeVarargs
+  private void putInternal(FullEntity<Key> entity) {
+    Key key = entity.key();
+    toAdd.remove(key);
+    toUpdate.remove(key);
+    toDelete.remove(key);
+    toPut.put(key, entity);
+  }
+
   @Override
-  public final void put(Entity... entities) {
+  public final Entity put(FullEntity<?> entity) {
+    return DatastoreHelper.put(this, entity);
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public final void putWithDeferredIdAllocation(FullEntity<?>... entities) {
     validateActive();
-    for (Entity entity : entities) {
-      Key key = entity.key();
-      toAdd.remove(key);
-      toUpdate.remove(key);
-      toDelete.remove(key);
-      toPut.put(key, entity);
+    for (FullEntity<?> entity : entities) {
+      IncompleteKey key = entity.key();
+      Preconditions.checkArgument(key != null, "Entity must have a key");
+      if (key instanceof Key) {
+        putInternal(Entity.convert((FullEntity<Key>) entity));
+      } else {
+        toAddAutoId.add((FullEntity<IncompleteKey>) entity);
+      }
+    }
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public final List<Entity> put(FullEntity<?>... entities) {
+    validateActive();
+    List<IncompleteKey> incompleteKeys = Lists.newArrayListWithExpectedSize(entities.length);
+    for (FullEntity<?> entity : entities) {
+      IncompleteKey key = entity.key();
+      Preconditions.checkArgument(key != null, "Entity must have a key");
+      if (!(key instanceof Key)) {
+        incompleteKeys.add(key);
+      }
     }
+    Iterator<Key> allocated;
+    if (!incompleteKeys.isEmpty()) {
+      IncompleteKey[] toAllocate = Iterables.toArray(incompleteKeys, IncompleteKey.class);
+      allocated = datastore().allocateId(toAllocate).iterator();
+    } else {
+      allocated = Collections.emptyIterator();
+    }
+    List<Entity> answer = Lists.newArrayListWithExpectedSize(entities.length);
+    for (FullEntity<?> entity : entities) {
+      if (entity.key() instanceof Key) {
+        putInternal((FullEntity<Key>) entity);
+        answer.add(Entity.convert((FullEntity<Key>) entity));
+      } else {
+        Entity entityWithAllocatedId = Entity.builder(allocated.next(), entity).build();
+        putInternal(entityWithAllocatedId);
+        answer.add(entityWithAllocatedId);
+      }
+    }
+    return answer;
   }
 
   @Override
diff --git a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/Datastore.java b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
index 8b2afcb5ad0c..efc56d532751 100644
--- a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
+++ b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/Datastore.java
@@ -77,28 +77,35 @@ interface TransactionCallable<T> {
    * @throws DatastoreException upon failure
    * @see #allocateId(IncompleteKey)
    */
-  List<Key> allocateId(IncompleteKey... key);
+  List<Key> allocateId(IncompleteKey... keys);
 
   /**
    * {@inheritDoc}
    * @throws DatastoreException upon failure
    */
   @Override
-  void update(Entity... entity);
+  void update(Entity... entities);
 
   /**
    * {@inheritDoc}
    * @throws DatastoreException upon failure
    */
   @Override
-  void put(Entity... entity);
+  Entity put(FullEntity<?> entity);
 
   /**
    * {@inheritDoc}
    * @throws DatastoreException upon failure
    */
   @Override
-  void delete(Key... key);
+  List<Entity> put(FullEntity<?>... entities);
+
+  /**
+   * {@inheritDoc}
+   * @throws DatastoreException upon failure
+   */
+  @Override
+  void delete(Key... keys);
 
   /**
    * Returns a new KeyFactory for this service
diff --git a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java
index 918711aa2822..1c067a719585 100644
--- a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java
+++ b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreBatchWriter.java
@@ -19,60 +19,72 @@
 import java.util.List;
 
 /**
- * An interface to represent a batch of write operations.
- * All write operation for a batch writer will be applied to the Datastore in one RPC call.
+ * An interface to represent a batch of write operations. All write operation for a batch writer
+ * will be applied to the Datastore in one RPC call.
  */
 interface DatastoreBatchWriter extends DatastoreWriter {
 
   /**
-   * Datastore add operation.
-   * This method will also allocate id for any entity with an incomplete key.
-   * As oppose to {@link #add(FullEntity)}, this method will defer any necessary id allocation
+   * Datastore add operation. This method will also allocate id for any entity with an incomplete
+   * key. As opposed to {@link #add(FullEntity)}, this method will defer any necessary id allocation
    * to submit time.
    *
    * @throws IllegalArgumentException if any of the given entities is missing a key
-   * @throws DatastoreException if a given entity with a
-   *     complete key was already added to this writer or if not active
+   * @throws DatastoreException if a given entity with a complete key was already added to this
+   *     writer or if not active
    */
-  void addWithDeferredIdAllocation(FullEntity<?>... entity);
+  void addWithDeferredIdAllocation(FullEntity<?>... entities);
 
   /**
    * {@inheritDoc}
    * For entities with complete keys that were marked for deletion in this writer the operation
    * will be changed to {@link #put}.
-   * @throws DatastoreException if a given entity with the
-   *     same complete key was already added to this writer, if writer is not active or
-   *     if id allocation for an entity with an incomplete key failed.
+   *
+   * @throws DatastoreException if a given entity with the same complete key was already added to
+   *     this writer, if writer is not active or if id allocation for an entity with an incomplete
+   *     key failed.
    */
   @Override
-  List<Entity> add(FullEntity<?>... entity);
+  List<Entity> add(FullEntity<?>... entities);
 
   /**
    * {@inheritDoc}
    * This operation will be converted to {@link #put} operation for entities that were already
-   *     added or put in this writer
-   * @throws DatastoreException if an entity is marked for
-   *     deletion in this writer or if not active
+   * added or put in this writer.
+   *
+   * @throws DatastoreException if an entity is marked for deletion in this writer or if not active
    */
   @Override
-  void update(Entity... entity);
+  void update(Entity... entities);
 
   /**
    * {@inheritDoc}
    * This operation will also remove from this batch any prior writes for entities with the same
-   *     keys
+   * keys
+   *
    * @throws DatastoreException if not active
    */
   @Override
-  void delete(Key... key);
+  void delete(Key... keys);
+
+  /**
+   * Datastore put operation. This method will also allocate id for any entity with an incomplete
+   * key. As opposed to {@link #put(FullEntity[])}, this method will defer any necessary id
+   * allocation to submit time.
+   *
+   * @throws IllegalArgumentException if any of the given entities is missing a key
+   * @throws DatastoreException if not active
+   */
+  void putWithDeferredIdAllocation(FullEntity<?>... entities);
 
   /**
    * {@inheritDoc}
    * This operation will also remove from this writer any prior writes for the same entities.
+   *
    * @throws DatastoreException if not active
    */
   @Override
-  void put(Entity... entity);
+  List<Entity> put(FullEntity<?>... entities);
 
   /**
    * Returns {@code true} if still active (write operations were not sent to the Datastore).
diff --git a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java
index 04b3e96c5175..655b1722b1ca 100644
--- a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java
+++ b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreHelper.java
@@ -51,6 +51,10 @@ static Entity add(DatastoreWriter writer, FullEntity<?> entity) {
     return writer.add(new FullEntity<?>[] {entity}).get(0);
   }
 
+  static Entity put(DatastoreWriter writer, FullEntity<?> entity) {
+    return writer.put(new FullEntity<?>[] {entity}).get(0);
+  }
+
   static KeyFactory newKeyFactory(DatastoreOptions options) {
     return new KeyFactory(options.projectId(), options.namespace());
   }
diff --git a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index 07ca76d2118f..69a6c72ce83f 100644
--- a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -87,11 +87,12 @@ com.google.datastore.v1beta3.RunQueryResponse runQuery(
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1beta3.RunQueryResponse>() {
-        @Override public com.google.datastore.v1beta3.RunQueryResponse call()
-            throws DatastoreException {
-          return datastoreRpc.runQuery(requestPb);
-        }
-      }, retryParams, EXCEPTION_HANDLER, options().clock());
+            @Override
+            public com.google.datastore.v1beta3.RunQueryResponse call()
+                throws DatastoreException {
+              return datastoreRpc.runQuery(requestPb);
+            }
+          }, retryParams, EXCEPTION_HANDLER, options().clock());
     } catch (RetryHelperException e) {
       throw DatastoreException.translateAndThrow(e);
     }
@@ -125,11 +126,12 @@ com.google.datastore.v1beta3.AllocateIdsResponse allocateIds(
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1beta3.AllocateIdsResponse>() {
-        @Override public com.google.datastore.v1beta3.AllocateIdsResponse call()
-            throws DatastoreException {
-          return datastoreRpc.allocateIds(requestPb);
-        }
-      }, retryParams, EXCEPTION_HANDLER, options().clock());
+            @Override
+            public com.google.datastore.v1beta3.AllocateIdsResponse call()
+                throws DatastoreException {
+              return datastoreRpc.allocateIds(requestPb);
+            }
+          }, retryParams, EXCEPTION_HANDLER, options().clock());
     } catch (RetryHelperException e) {
       throw DatastoreException.translateAndThrow(e);
     }
@@ -166,7 +168,7 @@ public List<Entity> add(FullEntity<?>... entities) {
             "Duplicate entity with the key %s", entity.key());
         }
       } else {
-        Preconditions.checkArgument(entity.hasKey(), "entity %s is missing a key", entity);
+        Preconditions.checkArgument(entity.hasKey(), "Entity %s is missing a key", entity);
       }
       mutationsPb.add(com.google.datastore.v1beta3.Mutation.newBuilder()
           .setInsert(entity.toPb()).build());
@@ -281,19 +283,19 @@ com.google.datastore.v1beta3.LookupResponse lookup(
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1beta3.LookupResponse>() {
-        @Override public com.google.datastore.v1beta3.LookupResponse call()
-            throws DatastoreException {
-          return datastoreRpc.lookup(requestPb);
-        }
-      }, retryParams, EXCEPTION_HANDLER, options().clock());
+            @Override
+            public com.google.datastore.v1beta3.LookupResponse call()
+                throws DatastoreException {
+              return datastoreRpc.lookup(requestPb);
+            }
+          }, retryParams, EXCEPTION_HANDLER, options().clock());
     } catch (RetryHelperException e) {
       throw DatastoreException.translateAndThrow(e);
     }
   }
 
-  @SafeVarargs
   @Override
-  public final void update(Entity... entities) {
+  public void update(Entity... entities) {
     if (entities.length > 0) {
       List<com.google.datastore.v1beta3.Mutation> mutationsPb =
           new ArrayList<>();
@@ -309,22 +311,47 @@ public final void update(Entity... entities) {
     }
   }
 
-  @SafeVarargs
   @Override
-  public final void put(Entity... entities) {
-    if (entities.length > 0) {
-      List<com.google.datastore.v1beta3.Mutation> mutationsPb =
-          new ArrayList<>();
-      Map<Key, Entity> dedupEntities = new LinkedHashMap<>();
-      for (Entity entity : entities) {
-        dedupEntities.put(entity.key(), entity);
-      }
-      for (Entity e : dedupEntities.values()) {
+  public Entity put(FullEntity<?> entity) {
+    return DatastoreHelper.put(this, entity);
+  }
+
+  @SuppressWarnings("unchecked")
+  @Override
+  public List<Entity> put(FullEntity<?>... entities) {
+    if (entities.length == 0) {
+      return Collections.emptyList();
+    }
+    List<com.google.datastore.v1beta3.Mutation> mutationsPb = new ArrayList<>();
+    Map<Key, Entity> dedupEntities = new LinkedHashMap<>();
+    for (FullEntity<?> entity : entities) {
+      Preconditions.checkArgument(entity.hasKey(), "Entity %s is missing a key", entity);
+      if (entity.key() instanceof Key) {
+        Entity completeEntity = Entity.convert((FullEntity<Key>) entity);
+        dedupEntities.put(completeEntity.key(), completeEntity);
+      } else {
         mutationsPb.add(
-            com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(e.toPb()).build());
+            com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(entity.toPb()).build());
+      }
+    }
+    for (Entity entity : dedupEntities.values()) {
+      mutationsPb.add(
+          com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(entity.toPb()).build());
+    }
+    com.google.datastore.v1beta3.CommitResponse commitResponse = commitMutation(mutationsPb);
+    Iterator<com.google.datastore.v1beta3.MutationResult> mutationResults =
+        commitResponse.getMutationResultsList().iterator();
+    ImmutableList.Builder<Entity> responseBuilder = ImmutableList.builder();
+    for (FullEntity<?> entity : entities) {
+      Entity completeEntity = dedupEntities.get(entity.key());
+      if (completeEntity != null) {
+        responseBuilder.add(completeEntity);
+      } else {
+        responseBuilder.add(
+            Entity.builder(Key.fromPb(mutationResults.next().getKey()), entity).build());
       }
-      commitMutation(mutationsPb);
     }
+    return responseBuilder.build();
   }
 
   @Override
diff --git a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
index e3989f29ef8b..156203103d03 100644
--- a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
+++ b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreReader.java
@@ -40,7 +40,7 @@ public interface DatastoreReader {
    * @throws DatastoreException upon failure
    * @see #get(Key)
    */
-  Iterator<Entity> get(Key... key);
+  Iterator<Entity> get(Key... keys);
 
   /**
    * Returns a list with a value for each given key (ordered by input). {@code null} values are
diff --git a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java
index 2771bb39ea60..0c66149c2e2b 100644
--- a/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java
+++ b/gcloud-java-datastore/src/main/java/com/google/cloud/datastore/DatastoreWriter.java
@@ -24,8 +24,7 @@
 public interface DatastoreWriter {
 
   /**
-   * Datastore add operation.
-   * This method will automatically allocate an id if necessary.
+   * Datastore add operation. This method will automatically allocate an id if necessary.
    *
    * @param entity the entity to add
    * @return an {@code Entity} with the same properties and a key that is either newly allocated
@@ -36,8 +35,8 @@ public interface DatastoreWriter {
   Entity add(FullEntity<?> entity);
 
   /**
-   * Datastore add operation.
-   * This method will automatically allocate id for any entity with an incomplete key.
+   * Datastore add operation. This method will automatically allocate id for any entity with an
+   * incomplete key.
    *
    * @return a list of {@code Entity} ordered by input with the same properties and a key that
    *     is either newly allocated or the same one if was already complete
@@ -45,23 +44,39 @@ public interface DatastoreWriter {
    * @throws IllegalArgumentException if any of the given entities is missing a key
    * @see #add(FullEntity)
    */
-  List<Entity> add(FullEntity<?>... entity);
+  List<Entity> add(FullEntity<?>... entities);
 
   /**
-   * A Datastore update operation.
-   * The operation will fail if an entity with the same key does not already exist.
+   * A Datastore update operation. The operation will fail if an entity with the same key does not
+   * already exist.
    */
-  void update(Entity... entity);
+  void update(Entity... entities);
 
   /**
-   * A Datastore put (a.k.a upsert) operation.
-   * The operation will add or modify the entities.
+   * A Datastore put (a.k.a upsert) operation. This method will automatically allocate an id if
+   * necessary.
+   *
+   * @param entity the entity to put
+   * @return an {@code Entity} with the same properties and a key that is either newly allocated
+   *     or the same one if key is already complete
+   * @throws DatastoreException upon failure
+   * @throws IllegalArgumentException if the given entity is missing a key
+   */
+  Entity put(FullEntity<?> entity);
+
+  /**
+   * A Datastore put (a.k.a upsert) operation. This method will automatically allocate id for any
+   * entity with an incomplete key.
+   *
+   * @return a list of updated or inserted {@code Entity}, ordered by input. Returned keys are
+   *     either newly allocated or the same one if was already complete.
+   * @throws DatastoreException upon failure
+   * @throws IllegalArgumentException if any of the given entities is missing a key
    */
-  void put(Entity... entity);
+  List<Entity> put(FullEntity<?>... entities);
 
   /**
-   * A datastore delete operation.
-   * It is OK request a deletion of a non-existing entity.
+   * A datastore delete operation. It is OK request a deletion of a non-existing entity.
    */
-  void delete(Key... key);
+  void delete(Key... keys);
 }
diff --git a/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/BaseDatastoreBatchWriterTest.java b/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/BaseDatastoreBatchWriterTest.java
index 6a30314082b7..9fc7fb0a6dc2 100644
--- a/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/BaseDatastoreBatchWriterTest.java
+++ b/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/BaseDatastoreBatchWriterTest.java
@@ -211,8 +211,44 @@ public void testPut() throws Exception {
     pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(ENTITY1.toPb()).build());
     pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(ENTITY2.toPb()).build());
     pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(ENTITY3.toPb()).build());
-    batchWriter.put(ENTITY1, ENTITY2);
-    batchWriter.put(ENTITY3);
+    List<Entity> putEntities = batchWriter.put(ENTITY1, ENTITY2);
+    Entity putEntity = batchWriter.put(ENTITY3);
+    assertEquals(ENTITY1, putEntities.get(0));
+    assertEquals(ENTITY2, putEntities.get(1));
+    assertEquals(ENTITY3, putEntity);
+    assertEquals(pbs, batchWriter.toMutationPbList());
+  }
+
+  @Test
+  public void testPutIncompleteKey() throws Exception {
+    List<com.google.datastore.v1beta3.Mutation> pbs = new LinkedList<>();
+    pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(ENTITY1.toPb()).build());
+    pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder()
+        .setUpsert(Entity.builder(KEY2, INCOMPLETE_ENTITY_1).build().toPb())
+        .build());
+    pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder()
+        .setUpsert(Entity.builder(KEY3, INCOMPLETE_ENTITY_2).build().toPb())
+        .build());
+    Entity putEntity = batchWriter.put(ENTITY1);
+    List<Entity> putEntities = batchWriter.put(INCOMPLETE_ENTITY_1, INCOMPLETE_ENTITY_2);
+    assertEquals(ENTITY1, putEntity);
+    assertEquals(Entity.builder(KEY2, INCOMPLETE_ENTITY_1).build(), putEntities.get(0));
+    assertEquals(Entity.builder(KEY3, INCOMPLETE_ENTITY_2).build(), putEntities.get(1));
+    assertEquals(pbs, batchWriter.toMutationPbList());
+  }
+
+  @Test
+  public void testPutWithDeferredAllocation() throws Exception {
+    List<com.google.datastore.v1beta3.Mutation> pbs = new LinkedList<>();
+    pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder()
+        .setInsert(INCOMPLETE_ENTITY_1.toPb())
+        .build());
+    pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder()
+        .setInsert(INCOMPLETE_ENTITY_2.toPb())
+        .build());
+    pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(ENTITY1.toPb()).build());
+    batchWriter.put(ENTITY1);
+    batchWriter.putWithDeferredIdAllocation(INCOMPLETE_ENTITY_1, INCOMPLETE_ENTITY_2);
     assertEquals(pbs, batchWriter.toMutationPbList());
   }
 
@@ -221,8 +257,10 @@ public void testPutAfterPut() throws Exception {
     Entity entity = Entity.builder(ENTITY1).set("foo", "bar").build();
     List<com.google.datastore.v1beta3.Mutation> pbs = new LinkedList<>();
     pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(entity.toPb()).build());
-    batchWriter.put(ENTITY1);
-    batchWriter.put(entity);
+    Entity putEntity1 = batchWriter.put(ENTITY1);
+    Entity putEntity2 = batchWriter.put(entity);
+    assertEquals(ENTITY1, putEntity1);
+    assertEquals(entity, putEntity2);
     assertEquals(pbs, batchWriter.toMutationPbList());
   }
 
@@ -242,7 +280,8 @@ public void testPutAfterUpdate() throws Exception {
     List<com.google.datastore.v1beta3.Mutation> pbs = new LinkedList<>();
     pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(entity.toPb()).build());
     batchWriter.update(ENTITY1);
-    batchWriter.put(entity);
+    Entity putEntity = batchWriter.put(entity);
+    assertEquals(entity, putEntity);
     assertEquals(pbs, batchWriter.toMutationPbList());
   }
 
@@ -252,7 +291,8 @@ public void testPutAfterDelete() throws Exception {
     List<com.google.datastore.v1beta3.Mutation> pbs = new LinkedList<>();
     pbs.add(com.google.datastore.v1beta3.Mutation.newBuilder().setUpsert(entity.toPb()).build());
     batchWriter.delete(KEY1);
-    batchWriter.put(entity);
+    Entity putEntity = batchWriter.put(entity);
+    assertEquals(entity, putEntity);
     assertEquals(pbs, batchWriter.toMutationPbList());
   }
 
diff --git a/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
index d73d6964e064..75a4d00883d5 100644
--- a/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
+++ b/gcloud-java-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java
@@ -903,21 +903,23 @@ public void testUpdate() {
 
   @Test
   public void testPut() {
-    Iterator<Entity> keys =
-        datastore.fetch(ENTITY1.key(), ENTITY2.key(), ENTITY3.key()).iterator();
-    assertEquals(ENTITY1, keys.next());
-    assertEquals(ENTITY2, keys.next());
-    assertNull(keys.next());
-    assertFalse(keys.hasNext());
+    Entity updatedEntity = Entity.builder(ENTITY1).set("new_property", 42L).build();
+    assertEquals(updatedEntity, datastore.put(updatedEntity));
+    assertEquals(updatedEntity, datastore.get(updatedEntity.key()));
 
     Entity entity2 = Entity.builder(ENTITY2).clear().set("bla", new NullValue()).build();
     assertNotEquals(ENTITY2, entity2);
-    datastore.put(ENTITY3, ENTITY1, entity2);
-    keys = datastore.fetch(ENTITY1.key(), ENTITY2.key(), ENTITY3.key()).iterator();
-    assertEquals(ENTITY1, keys.next());
-    assertEquals(entity2, keys.next());
-    assertEquals(ENTITY3, keys.next());
-    assertFalse(keys.hasNext());
+    List<Entity> entities = datastore.put(ENTITY1, entity2, ENTITY3, PARTIAL_ENTITY1);
+    assertEquals(ENTITY1, entities.get(0));
+    assertEquals(entity2, entities.get(1));
+    assertEquals(ENTITY3, entities.get(2));
+    assertEquals(PARTIAL_ENTITY1.properties(), entities.get(3).properties());
+    assertEquals(PARTIAL_ENTITY1.key().ancestors(), entities.get(3).key().ancestors());
+    assertEquals(ENTITY1, datastore.get(ENTITY1.key()));
+    assertEquals(entity2, datastore.get(entity2.key()));
+    assertEquals(ENTITY3, datastore.get(ENTITY3.key()));
+    Entity entity = datastore.get(entities.get(3).key());
+    assertEquals(entities.get(3), entity);
   }
 
   @Test