Skip to content

Commit 35c81ff

Browse files
author
Bingqin Zhou
committed
Fix unit tests.
1 parent 6517a06 commit 35c81ff

File tree

5 files changed

+107
-111
lines changed

5 files changed

+107
-111
lines changed

‎kcbq-connector/src/main/java/com/wepay/kafka/connect/bigquery/BigQuerySinkTask.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ public class BigQuerySinkTask extends SinkTask {
8989

9090
private final BigQuery testBigQuery;
9191
private final Storage testGcs;
92+
private final SchemaManager testSchemaManager;
9293

9394
private final UUID uuid = UUID.randomUUID();
9495
private ScheduledExecutorService gcsLoadExecutor;
@@ -100,6 +101,7 @@ public BigQuerySinkTask() {
100101
testBigQuery = null;
101102
schemaRetriever = null;
102103
testGcs = null;
104+
testSchemaManager = null;
103105
}
104106

105107
/**
@@ -110,10 +112,11 @@ public BigQuerySinkTask() {
110112
* @param testGcs {@link Storage} to use for testing (likely a mock)
111113
* @see BigQuerySinkTask#BigQuerySinkTask()
112114
*/
113-
public BigQuerySinkTask(BigQuery testBigQuery, SchemaRetriever schemaRetriever, Storage testGcs) {
115+
public BigQuerySinkTask(BigQuery testBigQuery, SchemaRetriever schemaRetriever, Storage testGcs, SchemaManager testSchemaManager) {
114116
this.testBigQuery = testBigQuery;
115117
this.schemaRetriever = schemaRetriever;
116118
this.testGcs = testGcs;
119+
this.testSchemaManager = testSchemaManager;
117120
}
118121

119122
@Override
@@ -256,6 +259,9 @@ private BigQuery getBigQuery() {
256259
}
257260

258261
private SchemaManager getSchemaManager(BigQuery bigQuery) {
262+
if (testSchemaManager != null) {
263+
return testSchemaManager;
264+
}
259265
schemaRetriever = config.getSchemaRetriever();
260266
SchemaConverter<com.google.cloud.bigquery.Schema> schemaConverter =
261267
config.getSchemaConverter();

‎kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/BigQuerySinkConnectorTest.java‎

Lines changed: 0 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -76,96 +76,6 @@ public static void initializePropertiesFactory() {
7676
propertiesFactory = new SinkConnectorPropertiesFactory();
7777
}
7878

79-
@Test
80-
public void testAutoCreateTables() {
81-
final String dataset = "scratch";
82-
final String existingTableTopic = "topic-with-existing-table";
83-
final String nonExistingTableTopic = "topic-without-existing-table";
84-
final TableId existingTable = TableId.of(dataset, "topic_with_existing_table");
85-
final TableId nonExistingTable = TableId.of(dataset, "topic_without_existing_table");
86-
87-
Map<String, String> properties = propertiesFactory.getProperties();
88-
properties.put(BigQuerySinkConfig.TABLE_CREATE_CONFIG, "true");
89-
properties.put(BigQuerySinkConfig.SCHEMA_RETRIEVER_CONFIG, MockSchemaRetriever.class.getName());
90-
properties.put(BigQuerySinkConfig.SANITIZE_TOPICS_CONFIG, "true");
91-
properties.put(BigQuerySinkConfig.DATASETS_CONFIG, String.format(".*=%s", dataset));
92-
properties.put(
93-
BigQuerySinkConfig.TOPICS_CONFIG,
94-
String.format("%s, %s", existingTableTopic, nonExistingTableTopic)
95-
);
96-
97-
BigQuery bigQuery = mock(BigQuery.class);
98-
Table fakeTable = mock(Table.class);
99-
when(bigQuery.getTable(existingTable)).thenReturn(fakeTable);
100-
when(bigQuery.getTable(nonExistingTable)).thenReturn(null);
101-
102-
SchemaManager schemaManager = mock(SchemaManager.class);
103-
104-
BigQuerySinkConnector testConnector = new BigQuerySinkConnector(bigQuery, schemaManager);
105-
testConnector.start(properties);
106-
107-
verify(bigQuery).getTable(existingTable);
108-
verify(bigQuery).getTable(nonExistingTable);
109-
verify(schemaManager, never()).createTable(existingTable, existingTableTopic);
110-
verify(schemaManager).createTable(nonExistingTable, nonExistingTableTopic);
111-
}
112-
113-
@Test
114-
public void testNonAutoCreateTables() {
115-
final String dataset = "scratch";
116-
final String[] topics = new String[] { "topic-one", "topicTwo", "TOPIC_THREE", "topic.four" };
117-
final String[] tables = new String[] { "topic_one", "topicTwo", "TOPIC_THREE", "topic_four" };
118-
119-
Map<String, String> properties = propertiesFactory.getProperties();
120-
properties.put(BigQuerySinkConfig.TABLE_CREATE_CONFIG, "false");
121-
properties.put(BigQuerySinkConfig.SANITIZE_TOPICS_CONFIG, "true");
122-
properties.put(BigQuerySinkConfig.DATASETS_CONFIG, String.format(".*=%s", dataset));
123-
properties.put(BigQuerySinkConfig.TOPICS_CONFIG, String.join(",", topics));
124-
125-
BigQuery bigQuery = mock(BigQuery.class);
126-
for (String table : tables) {
127-
Table fakeTable = mock(Table.class);
128-
when(bigQuery.getTable(TableId.of(dataset, table))).thenReturn(fakeTable);
129-
}
130-
131-
SchemaManager schemaManager = mock(SchemaManager.class);
132-
133-
BigQuerySinkConnector testConnector = new BigQuerySinkConnector(bigQuery);
134-
testConnector.start(properties);
135-
136-
verify(schemaManager, never()).createTable(any(TableId.class), any(String.class));
137-
138-
for (String table : tables) {
139-
verify(bigQuery).getTable(TableId.of(dataset, table));
140-
}
141-
}
142-
143-
@Test(expected = BigQueryConnectException.class)
144-
public void testNonAutoCreateTablesFailure() {
145-
final String dataset = "scratch";
146-
final String existingTableTopic = "topic-with-existing-table";
147-
final String nonExistingTableTopic = "topic-without-existing-table";
148-
final TableId existingTable = TableId.of(dataset, "topic_with_existing_table");
149-
final TableId nonExistingTable = TableId.of(dataset, "topic_without_existing_table");
150-
151-
Map<String, String> properties = propertiesFactory.getProperties();
152-
properties.put(BigQuerySinkConfig.TABLE_CREATE_CONFIG, "false");
153-
properties.put(BigQuerySinkConfig.SANITIZE_TOPICS_CONFIG, "true");
154-
properties.put(BigQuerySinkConfig.DATASETS_CONFIG, String.format(".*=%s", dataset));
155-
properties.put(
156-
BigQuerySinkConfig.TOPICS_CONFIG,
157-
String.format("%s, %s", existingTableTopic, nonExistingTableTopic)
158-
);
159-
160-
BigQuery bigQuery = mock(BigQuery.class);
161-
Table fakeTable = mock(Table.class);
162-
when(bigQuery.getTable(existingTable)).thenReturn(fakeTable);
163-
when(bigQuery.getTable(nonExistingTable)).thenReturn(null);
164-
165-
BigQuerySinkConnector testConnector = new BigQuerySinkConnector(bigQuery);
166-
testConnector.start(properties);
167-
}
168-
16979
@Test
17080
public void testTaskClass() {
17181
assertEquals(BigQuerySinkTask.class, new BigQuerySinkConnector().taskClass());

‎kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/BigQuerySinkTaskTest.java‎

Lines changed: 94 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -27,13 +27,15 @@
2727
import static org.mockito.Mockito.times;
2828
import static org.mockito.Mockito.verify;
2929
import static org.mockito.Mockito.when;
30+
import static org.mockito.Mockito.never;
3031

3132
import com.google.cloud.bigquery.BigQuery;
3233
import com.google.cloud.bigquery.BigQueryError;
3334
import com.google.cloud.bigquery.BigQueryException;
3435
import com.google.cloud.bigquery.InsertAllRequest;
3536
import com.google.cloud.bigquery.InsertAllResponse;
3637
import com.google.cloud.bigquery.TableId;
38+
import com.google.cloud.bigquery.Table;
3739
import com.google.cloud.storage.Storage;
3840

3941
import com.wepay.kafka.connect.bigquery.api.SchemaRetriever;
@@ -85,7 +87,7 @@ public void testSimplePut() {
8587
when(bigQuery.insertAll(anyObject())).thenReturn(insertAllResponse);
8688
when(insertAllResponse.hasErrors()).thenReturn(false);
8789

88-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
90+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
8991
testTask.initialize(sinkTaskContext);
9092
testTask.start(properties);
9193

@@ -112,7 +114,7 @@ public void testSimplePutWhenSchemaRetrieverIsNotNull() {
112114

113115
SchemaRetriever schemaRetriever = mock(SchemaRetriever.class);
114116

115-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, schemaRetriever, storage);
117+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, schemaRetriever, storage, null);
116118
testTask.initialize(sinkTaskContext);
117119
testTask.start(properties);
118120

@@ -123,13 +125,91 @@ public void testSimplePutWhenSchemaRetrieverIsNotNull() {
123125
any(String.class), any(Schema.class));
124126
}
125127

128+
@Test
129+
public void testAutoCreateTables() {
130+
final String dataset = "scratch";
131+
final String existingTableTopic = "topic-with-existing-table";
132+
final String nonExistingTableTopic = "topic-without-existing-table";
133+
final TableId existingTable = TableId.of(dataset, "topic_with_existing_table");
134+
final TableId nonExistingTable = TableId.of(dataset, "topic_without_existing_table");
135+
136+
Map<String, String> properties = propertiesFactory.getProperties();
137+
properties.put(BigQuerySinkConfig.TABLE_CREATE_CONFIG, "true");
138+
properties.put(BigQuerySinkConfig.SCHEMA_RETRIEVER_CONFIG, BigQuerySinkConnectorTest.MockSchemaRetriever.class.getName());
139+
properties.put(BigQuerySinkConfig.SANITIZE_TOPICS_CONFIG, "true");
140+
properties.put(BigQuerySinkConfig.DATASETS_CONFIG, String.format(".*=%s", dataset));
141+
properties.put(BigQuerySinkConfig.TOPICS_CONFIG, existingTableTopic);
142+
143+
BigQuery bigQuery = mock(BigQuery.class);
144+
Table fakeTable = mock(Table.class);
145+
when(bigQuery.getTable(existingTable)).thenReturn(fakeTable);
146+
when(bigQuery.getTable(nonExistingTable)).thenReturn(null);
147+
InsertAllResponse insertAllResponse = mock(InsertAllResponse.class);
148+
when(bigQuery.insertAll(anyObject())).thenReturn(insertAllResponse);
149+
when(insertAllResponse.hasErrors()).thenReturn(false);
150+
151+
Storage storage = mock(Storage.class);
152+
SchemaRetriever schemaRetriever = mock(SchemaRetriever.class);
153+
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
154+
SchemaManager schemaManager = mock(SchemaManager.class);
155+
156+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, schemaRetriever, storage, schemaManager);
157+
testTask.initialize(sinkTaskContext);
158+
testTask.start(properties);
159+
160+
testTask.put(Collections.singletonList(spoofSinkRecord(nonExistingTableTopic)));
161+
testTask.flush(Collections.emptyMap());
162+
163+
verify(schemaManager, never()).createTable(existingTable, existingTableTopic);
164+
verify(schemaManager).createTable(nonExistingTable, nonExistingTableTopic);
165+
}
166+
167+
@Test
168+
public void testNonAutoCreateTables() {
169+
final String dataset = "scratch";
170+
final String existingTableTopic = "topic-with-existing-table";
171+
final String nonExistingTableTopic = "topic-without-existing-table";
172+
final TableId existingTable = TableId.of(dataset, "topic_with_existing_table");
173+
final TableId nonExistingTable = TableId.of(dataset, "topic_without_existing_table");
174+
175+
Map<String, String> properties = propertiesFactory.getProperties();
176+
properties.put(BigQuerySinkConfig.SCHEMA_RETRIEVER_CONFIG, BigQuerySinkConnectorTest.MockSchemaRetriever.class.getName());
177+
properties.put(BigQuerySinkConfig.SANITIZE_TOPICS_CONFIG, "true");
178+
properties.put(BigQuerySinkConfig.DATASETS_CONFIG, String.format(".*=%s", dataset));
179+
properties.put(BigQuerySinkConfig.TOPICS_CONFIG, existingTableTopic);
180+
181+
BigQuery bigQuery = mock(BigQuery.class);
182+
Table fakeTable = mock(Table.class);
183+
when(bigQuery.getTable(existingTable)).thenReturn(fakeTable);
184+
when(bigQuery.getTable(nonExistingTable)).thenReturn(null);
185+
InsertAllResponse insertAllResponse = mock(InsertAllResponse.class);
186+
when(bigQuery.insertAll(anyObject())).thenReturn(insertAllResponse);
187+
when(insertAllResponse.hasErrors()).thenReturn(false);
188+
189+
Storage storage = mock(Storage.class);
190+
SchemaRetriever schemaRetriever = mock(SchemaRetriever.class);
191+
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
192+
SchemaManager schemaManager = mock(SchemaManager.class);
193+
194+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, schemaRetriever, storage, schemaManager);
195+
testTask.initialize(sinkTaskContext);
196+
testTask.start(properties);
197+
198+
testTask.put(Collections.singletonList(spoofSinkRecord(nonExistingTableTopic)));
199+
testTask.flush(Collections.emptyMap());
200+
201+
verify(schemaManager, never()).createTable(existingTable, existingTableTopic);
202+
verify(schemaManager, never()).createTable(nonExistingTable, existingTableTopic);
203+
}
204+
205+
126206
@Test
127207
public void testEmptyPut() {
128208
Map<String, String> properties = propertiesFactory.getProperties();
129209
BigQuery bigQuery = mock(BigQuery.class);
130210
Storage storage = mock(Storage.class);
131211

132-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
212+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
133213
testTask.start(properties);
134214

135215
testTask.put(Collections.emptyList());
@@ -148,7 +228,7 @@ public void testEmptyRecordPut() {
148228
BigQuery bigQuery = mock(BigQuery.class);
149229
Storage storage = mock(Storage.class);
150230

151-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
231+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
152232
testTask.start(properties);
153233

154234
SinkRecord emptyRecord = spoofSinkRecord(topic, simpleSchema, null);
@@ -175,7 +255,7 @@ public void testPutWhenPartitioningOnMessageTime() {
175255
when(bigQuery.insertAll(anyObject())).thenReturn(insertAllResponse);
176256
when(insertAllResponse.hasErrors()).thenReturn(false);
177257

178-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
258+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
179259
testTask.initialize(sinkTaskContext);
180260
testTask.start(properties);
181261

@@ -206,7 +286,7 @@ public void testPutWhenPartitioningOnMessageTimeWhenNoTimestampType() {
206286
when(bigQuery.insertAll(anyObject())).thenReturn(insertAllResponse);
207287
when(insertAllResponse.hasErrors()).thenReturn(false);
208288

209-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
289+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
210290
testTask.initialize(sinkTaskContext);
211291
testTask.start(properties);
212292

@@ -232,7 +312,7 @@ public void testBufferClearOnFlushError() {
232312
.thenThrow(new RuntimeException("This is a test"));
233313

234314
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
235-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
315+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
236316
testTask.initialize(sinkTaskContext);
237317
testTask.start(properties);
238318

@@ -253,7 +333,7 @@ public void testEmptyFlush() {
253333
Storage storage = mock(Storage.class);
254334

255335
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
256-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
336+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
257337
testTask.initialize(sinkTaskContext);
258338
testTask.start(properties);
259339

@@ -284,7 +364,7 @@ public void testBigQuery5XXRetry() {
284364

285365
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
286366

287-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
367+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
288368
testTask.initialize(sinkTaskContext);
289369
testTask.start(properties);
290370
testTask.put(Collections.singletonList(spoofSinkRecord(topic)));
@@ -318,7 +398,7 @@ public void testBigQuery403Retry() {
318398

319399
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
320400

321-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
401+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
322402
testTask.initialize(sinkTaskContext);
323403
testTask.start(properties);
324404
testTask.put(Collections.singletonList(spoofSinkRecord(topic)));
@@ -349,7 +429,7 @@ public void testBigQueryRetryExceeded() {
349429

350430
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
351431

352-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
432+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
353433
testTask.initialize(sinkTaskContext);
354434
testTask.start(properties);
355435
testTask.put(Collections.singletonList(spoofSinkRecord(topic)));
@@ -374,7 +454,7 @@ public void testInterruptedException() {
374454
when(bigQuery.insertAll(any(InsertAllRequest.class))).thenReturn(fakeResponse);
375455

376456
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
377-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
457+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
378458
testTask.initialize(sinkTaskContext);
379459
testTask.start(properties);
380460

@@ -393,7 +473,7 @@ public void testConfigException() {
393473
badProperties.remove(BigQuerySinkConfig.TOPICS_CONFIG);
394474

395475
BigQuerySinkTask testTask =
396-
new BigQuerySinkTask(mock(BigQuery.class), null, mock(Storage.class));
476+
new BigQuerySinkTask(mock(BigQuery.class), null, mock(Storage.class), null);
397477
testTask.start(badProperties);
398478
}
399479

@@ -421,7 +501,7 @@ public void testStop() {
421501
when(insertAllResponse.hasErrors()).thenReturn(false);
422502

423503
Storage storage = mock(Storage.class);
424-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
504+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
425505
testTask.initialize(sinkTaskContext);
426506
testTask.start(properties);
427507
testTask.put(Collections.singletonList(spoofSinkRecord(topic)));

‎kcbq-connector/src/test/java/com/wepay/kafka/connect/bigquery/write/row/BigQueryWriterTest.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ public void testBigQueryNoFailure() {
8585
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
8686

8787
Storage storage = mock(Storage.class);
88-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
88+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
8989
testTask.initialize(sinkTaskContext);
9090
testTask.start(properties);
9191
testTask.put(
@@ -133,7 +133,7 @@ public void testBigQueryPartialFailure() {
133133
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
134134

135135
Storage storage = mock(Storage.class);
136-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
136+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
137137
testTask.initialize(sinkTaskContext);
138138
testTask.start(properties);
139139
testTask.put(sinkRecordList);
@@ -182,7 +182,7 @@ public void testBigQueryCompleteFailure() {
182182
SinkTaskContext sinkTaskContext = mock(SinkTaskContext.class);
183183

184184
Storage storage = mock(Storage.class);
185-
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage);
185+
BigQuerySinkTask testTask = new BigQuerySinkTask(bigQuery, null, storage, null);
186186
testTask.initialize(sinkTaskContext);
187187
testTask.start(properties);
188188
testTask.put(sinkRecordList);

0 commit comments

Comments
 (0)