src/test/java/com/googlesource/gerrit/plugins/spannerrefdb/EmulatedSpannerRefDb.java

// Copyright (C) 2023 The Android Open Source Project
//
// 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.googlesource.gerrit.plugins.spannerrefdb;

import com.google.cloud.NoCredentials;
import com.google.cloud.spanner.Database;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.DatabaseId;
import com.google.cloud.spanner.Instance;
import com.google.cloud.spanner.InstanceAdminClient;
import com.google.cloud.spanner.InstanceConfigId;
import com.google.cloud.spanner.InstanceId;
import com.google.cloud.spanner.InstanceInfo;
import com.google.cloud.spanner.Spanner;
import com.google.cloud.spanner.SpannerOptions;
import java.util.Collections;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import org.testcontainers.containers.SpannerEmulatorContainer;
import org.testcontainers.utility.DockerImageName;

public class EmulatedSpannerRefDb extends SpannerTestSystem {
  public static final String PROJECT_ID = "test";
  public static final String SPANNER_INSTANCE_ID = "spanner-instance";
  public static final String SPANNER_DATABASE_ID = "global-refdb";

  private final SpannerEmulatorContainer container;

  private final Spanner spanner;
  private final SpannerOptions spannerOptions;
  private final InstanceAdminClient instanceAdminClient;
  private final Instance spannerInstance;
  private final Database spannerDatabase;
  private final DatabaseId dbId;
  private final DatabaseClient databaseClient;
  private final ScheduledExecutorService heartbeatExecutor;
  private final SpannerRefDatabase spannerRefDb;

  public EmulatedSpannerRefDb() throws Exception {
    container =
        new SpannerEmulatorContainer(
            DockerImageName.parse("gcr.io/cloud-spanner-emulator/emulator").withTag("1.5.9"));
    container.start();
    System.out.println(
        String.format(
            "Spanner emulator container started and is listening on %s",
            container.getEmulatorGrpcEndpoint()));

    spannerOptions = getEmulatorOptions();
    spanner = spannerOptions.getService();
    instanceAdminClient = spanner.getInstanceAdminClient();
    spannerInstance = createSpannerInstance();
    dbId = DatabaseId.of(spannerOptions.getProjectId(), SPANNER_INSTANCE_ID, SPANNER_DATABASE_ID);
    spannerDatabase =
        spannerInstance.createDatabase(SPANNER_DATABASE_ID, Collections.emptyList()).get();
    createSchema();
    databaseClient = createDatabaseClient();
    heartbeatExecutor = Executors.newScheduledThreadPool(2);
    Lock.Factory lockFactory =
        new Lock.Factory() {
          @Override
          public Lock create(String projectName, String refName) {
            return new Lock(databaseClient, "", heartbeatExecutor, projectName, refName);
          }
        };
    spannerRefDb = new SpannerRefDatabase(databaseClient, lockFactory);
  }

  @Override
  void reset() throws Exception {
    // do nothing, a new container is created for each instance of this class
  }

  @Override
  SpannerRefDatabase database() {
    return spannerRefDb;
  }

  @Override
  DatabaseClient dbClient() {
    return databaseClient;
  }

  @Override
  public void cleanup() {
    heartbeatExecutor.shutdownNow();
    try {
      heartbeatExecutor.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    }
    spannerDatabase.drop();
    instanceAdminClient.deleteInstance(SPANNER_INSTANCE_ID);
    spanner.close();
    container.stop();
    container.close();
    System.out.println("Spanner emulator container was stopped");
  }

  private void createSchema() {
    DatabaseSchemaCreator dsc =
        new DatabaseSchemaCreator(spannerOptions.getService().getDatabaseAdminClient(), dbId);
    dsc.start();
  }

  public SpannerEmulatorContainer getContainer() {
    return container;
  }

  public Database getSpannerDatabase() {
    return spannerDatabase;
  }

  public SpannerRefDatabase getSpannerRefDatabase() {
    return spannerRefDb;
  }

  public DatabaseClient getDatabaseClient() {
    return databaseClient;
  }

  private SpannerOptions getEmulatorOptions() {
    return SpannerOptions.newBuilder()
        .setEmulatorHost(container.getEmulatorGrpcEndpoint())
        .setCredentials(NoCredentials.getInstance())
        .setProjectId(PROJECT_ID)
        .build();
  }

  private Instance createSpannerInstance() throws InterruptedException, ExecutionException {
    InstanceConfigId instanceConfig = InstanceConfigId.of(PROJECT_ID, "emulator-config");
    InstanceId instanceId = InstanceId.of(PROJECT_ID, SPANNER_INSTANCE_ID);
    InstanceInfo instanceInfo =
        InstanceInfo.newBuilder(instanceId)
            .setInstanceConfigId(instanceConfig)
            .setNodeCount(1)
            .setDisplayName("test instance")
            .build();
    return instanceAdminClient.createInstance(instanceInfo).get();
  }

  private DatabaseClient createDatabaseClient() {
    return spannerOptions.getService().getDatabaseClient(dbId);
  }
}