Skip to content

Commit

Permalink
chore: fix broken unit tests.
Browse files Browse the repository at this point in the history
  • Loading branch information
arpan14 committed Feb 12, 2024
1 parent d7f70fa commit 07c9cd1
Show file tree
Hide file tree
Showing 5 changed files with 79 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,17 @@
import com.google.cloud.Timestamp;
import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.protobuf.ByteString;
import com.google.protobuf.util.Timestamps;
import com.google.spanner.v1.Session;
import com.google.spanner.v1.Transaction;
import io.opentelemetry.api.OpenTelemetry;
import java.io.IOException;
import java.util.Collections;
import java.util.Map;
import org.junit.Before;
Expand All @@ -59,6 +64,10 @@ public final class BatchClientImplTest {

@Mock private SpannerRpc gapicRpc;
@Mock private SpannerOptions spannerOptions;
@Mock private InstanceAdminStubSettings instanceAdminStubSettings;
@Mock private DatabaseAdminStubSettings databaseAdminStubSettings;
@Mock private DatabaseAdminStub databaseAdminStub;
@Mock private InstanceAdminStub instanceAdminStub;
@Captor private ArgumentCaptor<Map<SpannerRpc.Option, Object>> optionsCaptor;
@Mock private BatchTransactionId txnID;

Expand All @@ -72,7 +81,7 @@ public static void setupOpenTelemetry() {

@SuppressWarnings("unchecked")
@Before
public void setUp() {
public void setUp() throws IOException {
initMocks(this);
DatabaseId db = DatabaseId.of(DB_NAME);
when(spannerOptions.getNumChannels()).thenReturn(4);
Expand All @@ -86,6 +95,10 @@ public void setUp() {
GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
when(transportOptions.getExecutorFactory()).thenReturn(mock(ExecutorFactory.class));
when(spannerOptions.getTransportOptions()).thenReturn(transportOptions);
when(instanceAdminStubSettings.createStub()).thenReturn(instanceAdminStub);
when(databaseAdminStubSettings.createStub()).thenReturn(databaseAdminStub);
when(gapicRpc.getInstanceAdminStubSettings()).thenReturn(instanceAdminStubSettings);
when(gapicRpc.getDatabaseAdminStubSettings()).thenReturn(databaseAdminStubSettings);
@SuppressWarnings("resource")
SpannerImpl spanner = new SpannerImpl(gapicRpc, spannerOptions);
client = new BatchClientImpl(spanner.getSessionClient(db));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,10 @@
import com.google.cloud.Timestamp;
import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
import com.google.protobuf.ByteString;
Expand All @@ -52,6 +56,7 @@
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.context.Scope;
import java.io.IOException;
import java.text.ParseException;
import java.util.Calendar;
import java.util.Collections;
Expand All @@ -76,6 +81,10 @@
public class SessionImplTest {
@Mock private SpannerRpc rpc;
@Mock private SpannerOptions spannerOptions;
@Mock private InstanceAdminStubSettings instanceAdminStubSettings;
@Mock private DatabaseAdminStubSettings databaseAdminStubSettings;
@Mock private DatabaseAdminStub databaseAdminStub;
@Mock private InstanceAdminStub instanceAdminStub;
private com.google.cloud.spanner.Session session;
@Captor private ArgumentCaptor<Map<SpannerRpc.Option, Object>> optionsCaptor;
private Map<SpannerRpc.Option, Object> options;
Expand All @@ -88,7 +97,7 @@ public static void setupOpenTelemetry() {

@SuppressWarnings("unchecked")
@Before
public void setUp() {
public void setUp() throws IOException {
MockitoAnnotations.initMocks(this);
when(spannerOptions.getNumChannels()).thenReturn(4);
when(spannerOptions.getPrefetchChunks()).thenReturn(1);
Expand All @@ -101,6 +110,10 @@ public void setUp() {
when(spannerOptions.getTransportOptions()).thenReturn(transportOptions);
when(spannerOptions.getSessionPoolOptions()).thenReturn(mock(SessionPoolOptions.class));
when(spannerOptions.getOpenTelemetry()).thenReturn(OpenTelemetry.noop());
when(instanceAdminStubSettings.createStub()).thenReturn(instanceAdminStub);
when(databaseAdminStubSettings.createStub()).thenReturn(databaseAdminStub);
when(rpc.getInstanceAdminStubSettings()).thenReturn(instanceAdminStubSettings);
when(rpc.getDatabaseAdminStubSettings()).thenReturn(databaseAdminStubSettings);
@SuppressWarnings("resource")
SpannerImpl spanner = new SpannerImpl(rpc, spannerOptions);
String dbName = "projects/p1/instances/i1/databases/d1";
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,14 @@
import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.spanner.SpannerException.DoNotConstructDirectly;
import com.google.cloud.spanner.SpannerImpl.ClosedException;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
import io.opentelemetry.api.OpenTelemetry;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.Collections;
Expand All @@ -55,6 +60,10 @@
public class SpannerImplTest {
@Mock private SpannerRpc rpc;
@Mock private SpannerOptions spannerOptions;
@Mock private InstanceAdminStubSettings instanceAdminStubSettings;
@Mock private DatabaseAdminStubSettings databaseAdminStubSettings;
@Mock private DatabaseAdminStub databaseAdminStub;
@Mock private InstanceAdminStub instanceAdminStub;
private SpannerImpl impl;

@Captor ArgumentCaptor<Map<SpannerRpc.Option, Object>> options;
Expand All @@ -66,7 +75,7 @@ public static void setupOpenTelemetry() {
}

@Before
public void setUp() {
public void setUp() throws IOException {
MockitoAnnotations.initMocks(this);
when(spannerOptions.getNumChannels()).thenReturn(4);
when(spannerOptions.getDatabaseRole()).thenReturn("role");
Expand All @@ -75,6 +84,10 @@ public void setUp() {
when(spannerOptions.getClock()).thenReturn(NanoClock.getDefaultClock());
when(spannerOptions.getSessionLabels()).thenReturn(Collections.emptyMap());
when(spannerOptions.getOpenTelemetry()).thenReturn(OpenTelemetry.noop());
when(instanceAdminStubSettings.createStub()).thenReturn(instanceAdminStub);
when(databaseAdminStubSettings.createStub()).thenReturn(databaseAdminStub);
when(rpc.getInstanceAdminStubSettings()).thenReturn(instanceAdminStubSettings);
when(rpc.getDatabaseAdminStubSettings()).thenReturn(databaseAdminStubSettings);
impl = new SpannerImpl(rpc, spannerOptions);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,10 @@
import com.google.cloud.grpc.GrpcTransportOptions;
import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
import com.google.cloud.spanner.TransactionManager.TransactionState;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.protobuf.ByteString;
import com.google.protobuf.Empty;
Expand All @@ -47,6 +51,7 @@
import com.google.spanner.v1.Session;
import com.google.spanner.v1.Transaction;
import io.opentelemetry.api.OpenTelemetry;
import java.io.IOException;
import java.util.Collections;
import java.util.UUID;
import java.util.concurrent.Executors;
Expand Down Expand Up @@ -77,6 +82,10 @@ public void release(ScheduledExecutorService exec) {

@Mock private SessionImpl session;
@Mock TransactionRunnerImpl.TransactionContextImpl txn;
@Mock private InstanceAdminStubSettings instanceAdminStubSettings;
@Mock private DatabaseAdminStubSettings databaseAdminStubSettings;
@Mock private DatabaseAdminStub databaseAdminStub;
@Mock private InstanceAdminStub instanceAdminStub;
private TransactionManagerImpl manager;

@BeforeClass
Expand Down Expand Up @@ -201,7 +210,7 @@ public void commitAfterRollbackFails() {

@SuppressWarnings("unchecked")
@Test
public void usesPreparedTransaction() {
public void usesPreparedTransaction() throws IOException {
SpannerOptions options = mock(SpannerOptions.class);
when(options.getNumChannels()).thenReturn(4);
GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
Expand Down Expand Up @@ -248,6 +257,11 @@ public void usesPreparedTransaction() {
com.google.protobuf.Timestamp.newBuilder()
.setSeconds(System.currentTimeMillis() * 1000))
.build()));
when(instanceAdminStubSettings.createStub()).thenReturn(instanceAdminStub);
when(databaseAdminStubSettings.createStub()).thenReturn(databaseAdminStub);
when(rpc.getInstanceAdminStubSettings()).thenReturn(instanceAdminStubSettings);
when(rpc.getDatabaseAdminStubSettings()).thenReturn(databaseAdminStubSettings);

DatabaseId db = DatabaseId.of("test", "test", "test");
try (SpannerImpl spanner = new SpannerImpl(rpc, options)) {
DatabaseClient client = spanner.getDatabaseClient(db);
Expand All @@ -263,7 +277,7 @@ public void usesPreparedTransaction() {

@SuppressWarnings({"unchecked", "resource"})
@Test
public void inlineBegin() {
public void inlineBegin() throws IOException {
SpannerOptions options = mock(SpannerOptions.class);
when(options.getNumChannels()).thenReturn(4);
GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
Expand Down Expand Up @@ -332,6 +346,11 @@ public void inlineBegin() {
com.google.protobuf.Timestamp.newBuilder()
.setSeconds(System.currentTimeMillis() * 1000))
.build()));
when(instanceAdminStubSettings.createStub()).thenReturn(instanceAdminStub);
when(databaseAdminStubSettings.createStub()).thenReturn(databaseAdminStub);
when(rpc.getInstanceAdminStubSettings()).thenReturn(instanceAdminStubSettings);
when(rpc.getDatabaseAdminStubSettings()).thenReturn(databaseAdminStubSettings);

DatabaseId db = DatabaseId.of("test", "test", "test");
try (SpannerImpl spanner = new SpannerImpl(rpc, options)) {
DatabaseClient client = spanner.getDatabaseClient(db);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,10 @@
import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory;
import com.google.cloud.spanner.SessionClient.SessionId;
import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub;
import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub;
import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
import com.google.cloud.spanner.spi.v1.SpannerRpc;
import com.google.common.base.Preconditions;
import com.google.protobuf.ByteString;
Expand Down Expand Up @@ -63,6 +67,7 @@
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.trace.Span;
import io.opentelemetry.context.Scope;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.UUID;
Expand All @@ -81,6 +86,7 @@
/** Unit test for {@link com.google.cloud.spanner.TransactionRunnerImpl} */
@RunWith(JUnit4.class)
public class TransactionRunnerImplTest {

private static final class TestExecutorFactory
implements ExecutorFactory<ScheduledExecutorService> {
@Override
Expand All @@ -97,6 +103,10 @@ public void release(ScheduledExecutorService exec) {
@Mock private SpannerRpc rpc;
@Mock private SessionImpl session;
@Mock private TransactionRunnerImpl.TransactionContextImpl txn;
@Mock private InstanceAdminStubSettings instanceAdminStubSettings;
@Mock private DatabaseAdminStubSettings databaseAdminStubSettings;
@Mock private DatabaseAdminStub databaseAdminStub;
@Mock private InstanceAdminStub instanceAdminStub;
private TransactionRunnerImpl transactionRunner;
private boolean firstRun;
private boolean usedInlinedBegin;
Expand Down Expand Up @@ -151,7 +161,7 @@ public void setUp() {

@SuppressWarnings("unchecked")
@Test
public void usesPreparedTransaction() {
public void usesPreparedTransaction() throws IOException {
SpannerOptions options = mock(SpannerOptions.class);
when(options.getNumChannels()).thenReturn(4);
GrpcTransportOptions transportOptions = mock(GrpcTransportOptions.class);
Expand Down Expand Up @@ -196,6 +206,11 @@ public void usesPreparedTransaction() {
.setCommitTimestamp(
Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000))
.build()));
when(instanceAdminStubSettings.createStub()).thenReturn(instanceAdminStub);
when(databaseAdminStubSettings.createStub()).thenReturn(databaseAdminStub);
when(rpc.getInstanceAdminStubSettings()).thenReturn(instanceAdminStubSettings);
when(rpc.getDatabaseAdminStubSettings()).thenReturn(databaseAdminStubSettings);

DatabaseId db = DatabaseId.of("test", "test", "test");
try (SpannerImpl spanner = new SpannerImpl(rpc, options)) {
DatabaseClient client = spanner.getDatabaseClient(db);
Expand Down

0 comments on commit 07c9cd1

Please sign in to comment.