diff --git a/it/src/main/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManager.java b/it/src/main/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManager.java index 142a8c2745..6e35c2c44b 100644 --- a/it/src/main/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManager.java +++ b/it/src/main/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManager.java @@ -45,7 +45,11 @@ public DefaultDatastoreResourceManager(Builder builder) { this.namespace = builder.namespace; this.datastore = - DatastoreOptions.newBuilder().setCredentials(builder.credentials).build().getService(); + DatastoreOptions.newBuilder() + .setProjectId(builder.project) + .setCredentials(builder.credentials) + .build() + .getService(); this.keys = new HashSet<>(); } @@ -111,17 +115,20 @@ public void cleanupAll() { keys.clear(); } - public static Builder builder(String namespace) { + public static Builder builder(String project, String namespace) { + checkArgument(!Strings.isNullOrEmpty(project), "project can not be empty"); checkArgument(!Strings.isNullOrEmpty(namespace), "namespace can not be empty"); - return new Builder(namespace); + return new Builder(project, namespace); } public static final class Builder { + private final String project; private final String namespace; private Credentials credentials; - private Builder(String namespace) { + private Builder(String project, String namespace) { + this.project = project; this.namespace = namespace; } diff --git a/it/src/test/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManagerIT.java b/it/src/test/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManagerIT.java index e5479444ed..26f00fe0f6 100644 --- a/it/src/test/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManagerIT.java +++ b/it/src/test/java/com/google/cloud/teleport/it/datastore/DefaultDatastoreResourceManagerIT.java @@ -30,7 +30,8 @@ public class DefaultDatastoreResourceManagerIT { @Test public void testInsert() throws IOException { DefaultDatastoreResourceManager resourceManager = - DefaultDatastoreResourceManager.builder(DatastoreUtils.createTestId("testInsert")) + DefaultDatastoreResourceManager.builder( + TestProperties.project(), DatastoreUtils.createTestId("testInsert")) .credentials(TestProperties.credentials()) .build(); List entities = @@ -49,7 +50,8 @@ public void testInsert() throws IOException { @Test public void testInsertQuery() throws IOException { DefaultDatastoreResourceManager resourceManager = - DefaultDatastoreResourceManager.builder(DatastoreUtils.createTestId("testInsertQuery")) + DefaultDatastoreResourceManager.builder( + TestProperties.project(), DatastoreUtils.createTestId("testInsertQuery")) .credentials(TestProperties.buildCredentialsFromEnv()) .build(); @@ -71,7 +73,8 @@ public void testInsertQuery() throws IOException { @Test public void testInsertCleanUp() throws IOException { DefaultDatastoreResourceManager resourceManager = - DefaultDatastoreResourceManager.builder(DatastoreUtils.createTestId("testInsertCleanUp")) + DefaultDatastoreResourceManager.builder( + TestProperties.project(), DatastoreUtils.createTestId("testInsertCleanUp")) .credentials(TestProperties.buildCredentialsFromEnv()) .build(); resourceManager.insert( diff --git a/v1/src/test/java/com/google/cloud/teleport/templates/BigQueryToDatastoreIT.java b/v1/src/test/java/com/google/cloud/teleport/templates/BigQueryToDatastoreIT.java index 15ac0157ec..6424a9dcbb 100644 --- a/v1/src/test/java/com/google/cloud/teleport/templates/BigQueryToDatastoreIT.java +++ b/v1/src/test/java/com/google/cloud/teleport/templates/BigQueryToDatastoreIT.java @@ -75,7 +75,7 @@ public void setup() { .setCredentials(credentials) .build(); datastoreResourceManager = - DefaultDatastoreResourceManager.builder(testId).credentials(credentials).build(); + DefaultDatastoreResourceManager.builder(PROJECT, testId).credentials(credentials).build(); } @After diff --git a/v1/src/test/java/com/google/cloud/teleport/templates/TextToDatastoreIT.java b/v1/src/test/java/com/google/cloud/teleport/templates/TextToDatastoreIT.java index f528400005..cd8ba20e53 100644 --- a/v1/src/test/java/com/google/cloud/teleport/templates/TextToDatastoreIT.java +++ b/v1/src/test/java/com/google/cloud/teleport/templates/TextToDatastoreIT.java @@ -50,10 +50,10 @@ public final class TextToDatastoreIT extends TemplateTestBase { @Before public void setup() { - testId = PipelineUtils.createJobName(""); + testId = PipelineUtils.createJobName("test"); datastoreResourceManager = - DefaultDatastoreResourceManager.builder(testId).credentials(credentials).build(); + DefaultDatastoreResourceManager.builder(PROJECT, testId).credentials(credentials).build(); } @After