<commit_before>/*
 * Copyright 2021 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
 *
 *     https://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.appengine.api.datastore;

import static com.google.appengine.api.datastore.DatastoreServiceConfig.Builder.withImplicitTransactionManagementPolicy;
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;

import com.google.appengine.api.datastore.DatastoreAttributes.DatastoreType;
import com.google.apphosting.api.ApiProxy;
import com.google.apphosting.datastore.DatastoreV3Pb.CompiledCursor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.Futures;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
import com.google.datastore.v1.BeginTransactionResponse;
import com.google.datastore.v1.CommitRequest;
import com.google.datastore.v1.CommitResponse;
import com.google.datastore.v1.EntityResult;
import com.google.datastore.v1.Key.PathElement;
import com.google.datastore.v1.Key.PathElement.IdTypeCase;
import com.google.datastore.v1.LookupRequest;
import com.google.datastore.v1.LookupResponse;
import com.google.datastore.v1.Mutation;
import com.google.datastore.v1.MutationResult;
import com.google.datastore.v1.QueryResultBatch;
import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
import com.google.datastore.v1.RollbackRequest;
import com.google.datastore.v1.RollbackResponse;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.rpc.Code;
import com.google.storage.onestore.v3.OnestoreEntity.IndexPosition;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicLong;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.easymock.EasyMock;
import org.junit.After;
import org.junit.Before;

/**
 * Base class for Cloud Datastore v1 service tests. For ease of testing, the create.* and
 * expect.*Request methods in this class take care of populating the project ID in requests.
 */
public abstract class BaseCloudDatastoreV1ServiceImplTest {

  // Incrementing transaction handle to ensure sequential transactions are unique.
  private final AtomicLong transactionId = new AtomicLong(1L);

  // Incrementing cursor position to ensure they are unique.
  private final AtomicLong cursorIndex = new AtomicLong(1L);

  private List<Object> mocks;

  static final String APP_ID = "s~project-id";

  CloudDatastoreV1Client cloudDatastoreV1Client;
  TransactionStack txnStack;
  DatastoreServiceConfig datastoreServiceConfig;

  @Before
  public final void setUpBase() throws Exception {
    DatastoreServiceGlobalConfig.setConfig(
        DatastoreServiceGlobalConfig.builder().appId(APP_ID).build());

    cloudDatastoreV1Client = EasyMock.createMock(CloudDatastoreV1Client.class);

    txnStack = new TransactionStackImpl(new InstanceMemberThreadLocalTransactionStack());
    datastoreServiceConfig =
        withImplicitTransactionManagementPolicy(ImplicitTransactionManagementPolicy.AUTO)
            .readPolicy(new ReadPolicy(ReadPolicy.Consistency.STRONG));
    resetMocks();
  }

  @After
  public void tearDown() throws Exception {
    DatastoreServiceGlobalConfig.clear();
  }

  DatastoreServiceImpl newDatastoreService() {
    return new DatastoreServiceImpl(newAsyncDatastoreService());
  }

  AsyncCloudDatastoreV1ServiceImpl newAsyncDatastoreService() {
    return newAsyncDatastoreService(null);
  }

  AsyncCloudDatastoreV1ServiceImpl newAsyncDatastoreService(final DatastoreType type) {
    return newAsyncDatastoreService(datastoreServiceConfig, txnStack, type);
  }

  AsyncCloudDatastoreV1ServiceImpl newAsyncDatastoreService(
      DatastoreServiceConfig dsConfig, TransactionStack txnStack, final DatastoreType type) {
    dsConfig = new DatastoreServiceConfig(dsConfig);
    ApiProxy.ApiConfig apiConfig = new ApiProxy.ApiConfig();
    apiConfig.setDeadlineInSeconds(dsConfig.getDeadline());
    if (type == null) {
      return new AsyncCloudDatastoreV1ServiceImpl(dsConfig, cloudDatastoreV1Client, txnStack);
    } else {
      return new AsyncCloudDatastoreV1ServiceImpl(dsConfig, cloudDatastoreV1Client, txnStack) {
        @Override
        public DatastoreType getDatastoreType() {
          return type;
        }
      };
    }
  }

  Key createKey(int index) {
    return KeyFactory.createKey("testkind" + index, "testname" + index);
  }

  Key createChildKey(Key parent, int index) {
    return KeyFactory.createKey(parent, "testkind" + index, "testname" + index);
  }

  Key toKeyWithDifferentAppId(Key key, String newAppId) {
    return new Key(
        key.getKind(),
        key.getParent() == null ? null : toKeyWithDifferentAppId(key.getParent(), newAppId),
        key.getId(),
        key.getName(),
        new AppIdNamespace(newAppId, key.getAppIdNamespace().getNamespace()));
  }

  Entity createEntity(Key key, int index) {
    Entity result = new Entity(key);
    result.setProperty("testprop" + index, "testvalue" + index);
    return result;
  }

  void assertEntityEquals(Entity expected, Entity actual) {
    assertWithMessage("expected: %s; actual: %s")
        .that(expected.getKey().equals(actual.getKey(), false))
        .isTrue();
    assertThat(actual.getProperty("aString")).isEqualTo(expected.getProperty("aString"));

    // Datastore may convert from Integer to Long.
    Number expectedInt = (Number) expected.getProperty("anInteger");
    Number actualInt = (Number) actual.getProperty("anInteger");
    assertThat(actualInt.longValue()).isEqualTo(expectedInt.longValue());

    // Datastore may convert from Float to Double.
    Number expectedFloat = (Number) expected.getProperty("aFloat");
    Number actualFloat = (Number) actual.getProperty("aFloat");
    assertThat(actualFloat.doubleValue()).isWithin(1e-9).of(expectedFloat.doubleValue());
  }

  /**
   * Create a LookupResponse with the given Entities. These all represent "found" entities and
   * should not be null.
   */
  LookupResponse createLookupResponse(Entity... entities) {
    return createLookupResponse(
        Arrays.asList(entities), Collections.<Key>emptyList(), Collections.<Key>emptyList());
  }

  LookupResponse createLookupResponse(
      Collection<Entity> entities, Collection<Key> missingKeys, Collection<Key> deferredKeys) {
    LookupResponse.Builder responseBuilder = LookupResponse.newBuilder();
    for (Entity entity : entities) {
      responseBuilder.addFoundBuilder().setEntity(DataTypeTranslator.toV1Entity(entity));
    }
    for (Key missingKey : missingKeys) {<commit_msg>Internal change<commit_after>/*
 * Copyright 2021 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
 *
 *     https://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.appengine.api.datastore;

import static com.google.appengine.api.datastore.DatastoreServiceConfig.Builder.withImplicitTransactionManagementPolicy;
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;

import com.google.appengine.api.datastore.DatastoreAttributes.DatastoreType;
import com.google.apphosting.api.ApiProxy;
import com.google.apphosting.datastore.DatastoreV3Pb.CompiledCursor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.Futures;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
import com.google.datastore.v1.BeginTransactionResponse;
import com.google.datastore.v1.CommitRequest;
import com.google.datastore.v1.CommitResponse;
import com.google.datastore.v1.EntityResult;
import com.google.datastore.v1.Key.PathElement;
import com.google.datastore.v1.Key.PathElement.IdTypeCase;
import com.google.datastore.v1.LookupRequest;
import com.google.datastore.v1.LookupResponse;
import com.google.datastore.v1.Mutation;
import com.google.datastore.v1.MutationResult;
import com.google.datastore.v1.QueryResultBatch;
import com.google.datastore.v1.QueryResultBatch.MoreResultsType;
import com.google.datastore.v1.RollbackRequest;
import com.google.datastore.v1.RollbackResponse;
import com.google.datastore.v1.RunQueryRequest;
import com.google.datastore.v1.RunQueryResponse;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.rpc.Code;
import com.google.storage.onestore.v3.OnestoreEntity.IndexPosition;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicLong;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.easymock.EasyMock;
import org.junit.After;
import org.junit.Before;

/**
 * Base class for Cloud Datastore v1 service tests. For ease of testing, the create.* and
 * expect.*Request methods in this class take care of populating the project ID in requests.
 */
public abstract class BaseCloudDatastoreV1ServiceImplTest {

  // Incrementing transaction handle to ensure sequential transactions are unique.
  private final AtomicLong transactionId = new AtomicLong(1L);

  // Incrementing cursor position to ensure they are unique.
  private final AtomicLong cursorIndex = new AtomicLong(1L);

  private List<Object> mocks;

  static final String APP_ID = "s~project-id";

  CloudDatastoreV1Client cloudDatastoreV1Client;
  TransactionStack txnStack;
  DatastoreServiceConfig datastoreServiceConfig;

  @Before
  public final void setUpBase() throws Exception {
    DatastoreServiceGlobalConfig.setConfig(
        DatastoreServiceGlobalConfig.builder().appId(APP_ID).build());

    cloudDatastoreV1Client = EasyMock.createMock(CloudDatastoreV1Client.class);

    txnStack = new TransactionStackImpl(new InstanceMemberThreadLocalTransactionStack());
    datastoreServiceConfig =
        withImplicitTransactionManagementPolicy(ImplicitTransactionManagementPolicy.AUTO)
            .readPolicy(new ReadPolicy(ReadPolicy.Consistency.STRONG));
    resetMocks();
  }

  @After
  public void tearDown() throws Exception {
    DatastoreServiceGlobalConfig.clear();
  }

  DatastoreServiceImpl newDatastoreService() {
    return new DatastoreServiceImpl(newAsyncDatastoreService());
  }

  AsyncCloudDatastoreV1ServiceImpl newAsyncDatastoreService() {
    return newAsyncDatastoreService(null);
  }

  AsyncCloudDatastoreV1ServiceImpl newAsyncDatastoreService(final DatastoreType type) {
    return newAsyncDatastoreService(datastoreServiceConfig, txnStack, type);
  }

  AsyncCloudDatastoreV1ServiceImpl newAsyncDatastoreService(
      DatastoreServiceConfig dsConfig, TransactionStack txnStack, final DatastoreType type) {
    dsConfig = new DatastoreServiceConfig(dsConfig);
    ApiProxy.ApiConfig apiConfig = new ApiProxy.ApiConfig();
    apiConfig.setDeadlineInSeconds(dsConfig.getDeadline());
    if (type == null) {
      return new AsyncCloudDatastoreV1ServiceImpl(dsConfig, cloudDatastoreV1Client, txnStack);
    } else {
      return new AsyncCloudDatastoreV1ServiceImpl(dsConfig, cloudDatastoreV1Client, txnStack) {
        @Override
        public DatastoreType getDatastoreType() {
          return type;
        }
      };
    }
  }

  Key createKey(int index) {
    return KeyFactory.createKey("testkind" + index, "testname" + index);
  }

  Key createChildKey(Key parent, int index) {
    return KeyFactory.createKey(parent, "testkind" + index, "testname" + index);
  }

  Key toKeyWithDifferentAppId(Key key, String newAppId) {
    return new Key(
        key.getKind(),
        key.getParent() == null ? null : toKeyWithDifferentAppId(key.getParent(), newAppId),
        key.getId(),
        key.getName(),
        new AppIdNamespace(newAppId, key.getAppIdNamespace().getNamespace()));
  }

  Entity createEntity(Key key, int index) {
    Entity result = new Entity(key);
    result.setProperty("testprop" + index, "testvalue" + index);
    return result;
  }

  void assertEntityEquals(Entity expected, Entity actual) {
    assertWithMessage("expected: %s; actual: %s", expected.getKey(), actual.getKey())
        .that(expected.getKey().equals(actual.getKey(), false))
        .isTrue();
    assertThat(actual.getProperty("aString")).isEqualTo(expected.getProperty("aString"));

    // Datastore may convert from Integer to Long.
    Number expectedInt = (Number) expected.getProperty("anInteger");
    Number actualInt = (Number) actual.getProperty("anInteger");
    assertThat(actualInt.longValue()).isEqualTo(expectedInt.longValue());

    // Datastore may convert from Float to Double.
    Number expectedFloat = (Number) expected.getProperty("aFloat");
    Number actualFloat = (Number) actual.getProperty("aFloat");
    assertThat(actualFloat.doubleValue()).isWithin(1e-9).of(expectedFloat.doubleValue());
  }

  /**
   * Create a LookupResponse with the given Entities. These all represent "found" entities and
   * should not be null.
   */
  LookupResponse createLookupResponse(Entity... entities) {
    return createLookupResponse(
        Arrays.asList(entities), Collections.<Key>emptyList(), Collections.<Key>emptyList());
  }

  LookupResponse createLookupResponse(
      Collection<Entity> entities, Collection<Key> missingKeys, Collection<Key> deferredKeys) {
    LookupResponse.Builder responseBuilder = LookupResponse.newBuilder();
    for (Entity entity : entities) {
      responseBuilder.addFoundBuilder().setEntity(DataTypeTranslator.toV1Entity(entity));
    }
    for (Key missingKey : missingKeys) {
{'fraction_non_alphanumeric': 0.09611829944547134, 'fraction_numerical': 0.006205439662001585, 'mean_word_length': 4.926838810641628, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1924248', 'n_tokens_mistral': 4424, 'n_tokens_neox': 4373, 'n_words': 959}