Add helper context method to fix test

Author: misolt

dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/teradata/TeradataLogsConnectorTest.java

@@ -18,7 +18,10 @@
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.Iterables.getOnlyElement;
+import static com.google.edwmigration.dumper.application.dumper.connector.teradata.AbstractTeradataConnector.DEF_LOG_TABLE;
+import static com.google.edwmigration.dumper.application.dumper.connector.teradata.AbstractTeradataConnector.DEF_SQL_TABLE;
 import static com.google.edwmigration.dumper.application.dumper.test.DumperTestUtils.assertQueryEquals;
+import static java.nio.file.FileSystems.newFileSystem;
 import static java.time.ZoneOffset.UTC;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -28,6 +31,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.edwmigration.dumper.application.dumper.ConnectorArguments;
 import com.google.edwmigration.dumper.application.dumper.connector.AbstractConnectorExecutionTest;
+import com.google.edwmigration.dumper.application.dumper.handle.Handle;
 import com.google.edwmigration.dumper.application.dumper.handle.JdbcHandle;
 import com.google.edwmigration.dumper.application.dumper.io.FileSystemOutputHandleFactory;
 import com.google.edwmigration.dumper.application.dumper.io.OutputHandleFactory;
@@ -37,7 +41,6 @@
 import java.io.File;
 import java.net.URI;
 import java.nio.file.FileSystem;
-import java.nio.file.FileSystems;
 import java.time.Clock;
 import java.time.Instant;
 import java.util.ArrayList;
@@ -53,7 +56,8 @@
 @RunWith(JUnit4.class)
 public class TeradataLogsConnectorTest extends AbstractConnectorExecutionTest {
 
-  final TaskRunContext mockContext = mock(TaskRunContext.class);
+  final TaskRunContextOps mockOptions = mock(TaskRunContextOps.class);
+
   private final TeradataLogsConnector connector = new TeradataLogsConnector();
 
   @Test
@@ -224,37 +228,30 @@ public void addTasksTo_commonConnectorTest_success() throws Exception {
 
   @Test
   public void addTasksTo_executeQuery_success() throws Exception {
-    String name = getClass().getSimpleName();
-    File dbFile = DumperTestUtils.newJdbcFile(name);
-    File zipFile = DumperTestUtils.newZipFile(name);
+    File dbFile = DumperTestUtils.newJdbcFile(getClass().getSimpleName());
+    File zipFile = DumperTestUtils.newZipFile(getClass().getSimpleName());
 
     dbFile.delete();
     zipFile.delete();
 
     URI outputUri = URI.create("jar:" + zipFile.toURI());
-
+    ImmutableList<String> scripts =
+        ImmutableList.of(
+            "attach ':memory:' as dbc",
+            String.format(
+                "create table %s (UserName varchar, errorcode int, StartTime int)", DEF_LOG_TABLE),
+            String.format("create table %s (QueryID int)", DEF_SQL_TABLE));
     // This isn't great because all the column-validity queries fail.
     try (JdbcHandle handle = DumperTestUtils.newJdbcHandle(dbFile);
-        FileSystem fileSystem =
-            FileSystems.newFileSystem(outputUri, ImmutableMap.of("create", "true"))) {
-      OutputHandleFactory sinkFactory = new FileSystemOutputHandleFactory(fileSystem, "/");
-      handle.getJdbcTemplate().execute("attach ':memory:' as dbc");
-      handle
-          .getJdbcTemplate()
-          .execute(
-              "create table "
-                  + TeradataLogsConnector.DEF_LOG_TABLE
-                  + " (UserName varchar, errorcode int, StartTime int)");
-      handle
-          .getJdbcTemplate()
-          .execute("create table " + TeradataLogsConnector.DEF_SQL_TABLE + " (QueryID int)");
-
+        FileSystem fileSystem = newFileSystem(outputUri, ImmutableMap.of("create", "true"))) {
+      scripts.forEach(handle.getJdbcTemplate()::execute);
       ConnectorArguments arguments =
-          new ConnectorArguments("--connector", connector.getName(), "--query-log-days", "1");
+          new ConnectorArguments("--connector", "teradata-logs", "--query-log-days", "1");
+      TaskRunContext runContext = testContext(fileSystem, handle, mockOptions, arguments);
       List<Task<?>> tasks = new ArrayList<>();
       connector.addTasksTo(tasks, arguments);
       for (Task<?> task : tasks) {
-        task.run(mockContext);
+        task.run(runContext);
       }
     }
   }
@@ -452,4 +449,13 @@ public void values_success() {
         Teradata14LogsConnector.EXPRESSIONS_LOG_TBL.toArray(
             new String[Teradata14LogsConnector.EXPRESSIONS_LOG_TBL.size()]));
   }
+
+  static TaskRunContext testContext(
+      FileSystem fileSystem,
+      Handle handle,
+      TaskRunContextOps options,
+      ConnectorArguments arguments) {
+    OutputHandleFactory sinkFactory = new FileSystemOutputHandleFactory(fileSystem, "/");
+    return new TaskRunContext(sinkFactory, handle, 10, options, arguments);
+  }
 }

dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/task/DumpMetadataTaskTest.java

@@ -35,8 +35,7 @@ public class DumpMetadataTaskTest extends AbstractTaskTest {
   public void testTask() throws Exception {
     MemoryByteSink sink = new MemoryByteSink();
     ConnectorArguments arguments = new ConnectorArguments("--connector", "bigquery-logs");
-    new DumpMetadataTask(arguments, "test-format")
-        .doRun(mockContext, sink, HANDLE);
+    new DumpMetadataTask(arguments, "test-format").doRun(mockContext, sink, HANDLE);
   }
 
   @Test

dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/task/JdbcSelectTaskTest.java

@@ -69,8 +69,7 @@ public static void setUpClass() throws Exception {
   public void testResultSetHeader() throws Exception {
     MemoryByteSink sink = new MemoryByteSink();
     try (JdbcHandle handle = DumperTestUtils.newJdbcHandle(FILE)) {
-      new JdbcSelectTask("(memory)", QUERY)
-          .doRun(mockContext, sink, handle);
+      new JdbcSelectTask("(memory)", QUERY).doRun(mockContext, sink, handle);
     }
     String actualOutput = sink.openStream().toString();
     assertEquals("a,b,c\n1,2,3\n", actualOutput);