Skip to content
This repository has been archived by the owner on Sep 27, 2023. It is now read-only.

refactor(samples): update run history it #348

Merged
merged 1 commit into from
Sep 18, 2020
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand All @@ -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<String, Value> 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);
Expand All @@ -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 :");
}
}