diff --git a/samples/snippets/src/test/java/com/example/bigquerydatatransfer/RunHistoryIT.java b/samples/snippets/src/test/java/com/example/bigquerydatatransfer/RunHistoryIT.java index 3bae3734..3324c00e 100644 --- a/samples/snippets/src/test/java/com/example/bigquerydatatransfer/RunHistoryIT.java +++ b/samples/snippets/src/test/java/com/example/bigquerydatatransfer/RunHistoryIT.java @@ -19,21 +19,9 @@ import static com.google.common.truth.Truth.assertThat; import static junit.framework.TestCase.assertNotNull; -import com.google.cloud.bigquery.BigQuery; -import com.google.cloud.bigquery.BigQueryOptions; -import com.google.cloud.bigquery.DatasetInfo; -import com.google.cloud.bigquery.datatransfer.v1.CreateTransferConfigRequest; -import com.google.cloud.bigquery.datatransfer.v1.DataTransferServiceClient; -import com.google.cloud.bigquery.datatransfer.v1.ProjectName; -import com.google.cloud.bigquery.datatransfer.v1.TransferConfig; -import com.google.protobuf.Struct; -import com.google.protobuf.Value; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.PrintStream; -import java.util.HashMap; -import java.util.Map; -import java.util.UUID; import java.util.logging.Level; import java.util.logging.Logger; import org.junit.After; @@ -44,16 +32,11 @@ public class RunHistoryIT { private static final Logger LOG = Logger.getLogger(GetTransferConfigInfoIT.class.getName()); - private static final String ID = UUID.randomUUID().toString().substring(0, 8); - private BigQuery bigquery; private ByteArrayOutputStream bout; - private String name; - private String displayName; - private String datasetName; private PrintStream out; private PrintStream originalPrintStream; - private static final String PROJECT_ID = requireEnvVar("GOOGLE_CLOUD_PROJECT"); + private static final String CONFIG_NAME = requireEnvVar("DTS_TRANSFER_CONFIG_NAME"); private static String requireEnvVar(String varName) { String value = System.getenv(varName); @@ -65,61 +48,19 @@ private static String requireEnvVar(String varName) { @BeforeClass public static void checkRequirements() { - requireEnvVar("GOOGLE_CLOUD_PROJECT"); + requireEnvVar("DTS_TRANSFER_CONFIG_NAME"); } @Before - public void setUp() throws IOException { + public void setUp() { bout = new ByteArrayOutputStream(); out = new PrintStream(bout); originalPrintStream = System.out; System.setOut(out); - - displayName = "MY_SCHEDULE_NAME_TEST_" + ID; - datasetName = "MY_DATASET_NAME_TEST_" + ID; - // create a temporary dataset - bigquery = BigQueryOptions.getDefaultInstance().getService(); - bigquery.create(DatasetInfo.of(datasetName)); - // create a scheduled query - String query = - "SELECT CURRENT_TIMESTAMP() as current_time, @run_time as intended_run_time, " - + "@run_date as intended_run_date, 17 as some_integer"; - String destinationTableName = - "MY_DESTINATION_TABLE_" + UUID.randomUUID().toString().substring(0, 8) + "_{run_date}"; - Map params = new HashMap<>(); - params.put("query", Value.newBuilder().setStringValue(query).build()); - params.put( - "destination_table_name_template", - Value.newBuilder().setStringValue(destinationTableName).build()); - params.put("write_disposition", Value.newBuilder().setStringValue("WRITE_TRUNCATE").build()); - params.put("partitioning_field", Value.newBuilder().setStringValue("").build()); - TransferConfig transferConfig = - TransferConfig.newBuilder() - .setDestinationDatasetId(datasetName) - .setDisplayName(displayName) - .setDataSourceId("scheduled_query") - .setParams(Struct.newBuilder().putAllFields(params).build()) - .setSchedule("every 24 hours") - .build(); - try (DataTransferServiceClient dataTransferServiceClient = DataTransferServiceClient.create()) { - ProjectName parent = ProjectName.of(PROJECT_ID); - CreateTransferConfigRequest request = - CreateTransferConfigRequest.newBuilder() - .setParent(parent.toString()) - .setTransferConfig(transferConfig) - .build(); - name = dataTransferServiceClient.createTransferConfig(request).getName(); - System.out.println("\nScheduled query created successfully :" + name); - } } @After - public void tearDown() throws IOException { - // delete scheduled query that was just created - DeleteScheduledQuery.deleteScheduledQuery(name); - // delete a temporary dataset - bigquery.delete(datasetName, BigQuery.DatasetDeleteOption.deleteContents()); - + public void tearDown() { // restores print statements in the original method System.out.flush(); System.setOut(originalPrintStream); @@ -128,7 +69,7 @@ public void tearDown() throws IOException { @Test public void testRunHistory() throws IOException { - RunHistory.runHistory(name); + RunHistory.runHistory(CONFIG_NAME); assertThat(bout.toString()).contains("Success! Run ID :"); } }