@Test
public void testSessionSave() throws Exception {

  DatastoreSession session = spy(new DatastoreSession(manager));
  session.setValid(true);
  session.setId(keyId);
  session.setAttribute("count", 5);
  store.save(session);
  ArgumentCaptor captor = ArgumentCaptor.forClass(Entity.class);
  verify(datastore).put((FullEntity<?>[]) captor.capture());
  verify(session).saveAttributesToEntity(any());
  List<Entity> entities = captor.getAllValues();
  assertEquals(2, entities.size());
  assertTrue(entities.stream()
      .map(e -> e.getKey().getName())
      .collect(Collectors.toList())
      .containsAll(Arrays.asList("count", keyId)));
}