2626import com .google .cloud .datastore .AggregationResult ;
2727import com .google .cloud .datastore .Datastore ;
2828import com .google .cloud .datastore .Datastore .TransactionCallable ;
29- import com .google .cloud .datastore .DatastoreOptions ;
3029import com .google .cloud .datastore .Entity ;
3130import com .google .cloud .datastore .EntityQuery ;
3231import com .google .cloud .datastore .GqlQuery ;
3332import com .google .cloud .datastore .Key ;
3433import com .google .cloud .datastore .Query ;
3534import com .google .cloud .datastore .QueryResults ;
3635import com .google .cloud .datastore .Transaction ;
37- import com .google .cloud .datastore .testing .RemoteDatastoreHelper ;
3836import com .google .common .collect .ImmutableList ;
3937import com .google .datastore .v1 .TransactionOptions ;
4038import com .google .datastore .v1 .TransactionOptions .ReadOnly ;
4341import java .util .concurrent .Executors ;
4442import java .util .concurrent .Future ;
4543import org .junit .After ;
46- import org .junit .AfterClass ;
44+ import org .junit .BeforeClass ;
45+ import org .junit .ClassRule ;
4746import org .junit .Test ;
4847
4948// TODO(jainsahab) Move all the aggregation related tests from ITDatastoreTest to this file
5049public class ITDatastoreAggregationsTest {
5150
52- private static final RemoteDatastoreHelper HELPER = RemoteDatastoreHelper . create ();
53- private static final DatastoreOptions OPTIONS = HELPER . getOptions ();
54- private static final Datastore DATASTORE = OPTIONS . getService () ;
51+ @ ClassRule public static MultiDbRule multiDbRule = new MultiDbRule ();
52+
53+ private static Datastore DATASTORE ;
5554
5655 private static final String KIND = "Marks" ;
5756
57+ @ BeforeClass
58+ public static void beforeClass () throws Exception {
59+ DATASTORE = multiDbRule .getDatastore ();
60+ }
61+
5862 @ After
5963 public void tearDown () {
6064 EntityQuery allEntitiesQuery = Query .newEntityQueryBuilder ().build ();
@@ -64,11 +68,6 @@ public void tearDown() {
6468 DATASTORE .delete (keysToDelete );
6569 }
6670
67- @ AfterClass
68- public static void afterClass () throws Exception {
69- DATASTORE .close ();
70- }
71-
7271 Key key1 = DATASTORE .newKeyFactory ().setKind (KIND ).newKey (1 );
7372 Key key2 = DATASTORE .newKeyFactory ().setKind (KIND ).newKey (2 );
7473 Key key3 = DATASTORE .newKeyFactory ().setKind (KIND ).newKey (3 );
@@ -89,7 +88,6 @@ public void testSumAggregation() {
8988 Query .newAggregationQueryBuilder ()
9089 .over (baseQuery )
9190 .addAggregations (sum ("marks" ).as ("total_marks" ))
92- .setNamespace (OPTIONS .getNamespace ())
9391 .build ();
9492
9593 // sum of 2 entities
@@ -108,11 +106,7 @@ public void testSumAggregationWithAutoGeneratedAlias() {
108106
109107 EntityQuery baseQuery = Query .newEntityQueryBuilder ().setKind (KIND ).build ();
110108 AggregationQuery aggregationQuery =
111- Query .newAggregationQueryBuilder ()
112- .over (baseQuery )
113- .addAggregations (sum ("marks" ))
114- .setNamespace (OPTIONS .getNamespace ())
115- .build ();
109+ Query .newAggregationQueryBuilder ().over (baseQuery ).addAggregations (sum ("marks" )).build ();
116110
117111 // sum of 2 entities
118112 assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getLong ("property_1" ))
@@ -133,11 +127,7 @@ public void testSumAggregationInGqlQuery() {
133127 "AGGREGATE SUM(marks) AS total_marks OVER (SELECT * FROM Marks)" )
134128 .build ();
135129
136- AggregationQuery aggregationQuery =
137- Query .newAggregationQueryBuilder ()
138- .over (gqlQuery )
139- .setNamespace (OPTIONS .getNamespace ())
140- .build ();
130+ AggregationQuery aggregationQuery = Query .newAggregationQueryBuilder ().over (gqlQuery ).build ();
141131
142132 // sum of 2 entities
143133 assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getLong ("total_marks" ))
@@ -158,7 +148,6 @@ public void testSumAggregationWithResultOfDoubleType() {
158148 Query .newAggregationQueryBuilder ()
159149 .over (baseQuery )
160150 .addAggregations (sum ("cgpa" ).as ("total_cgpa" ))
161- .setNamespace (OPTIONS .getNamespace ())
162151 .build ();
163152
164153 // sum of 2 entities
@@ -180,7 +169,6 @@ public void testAvgAggregation() {
180169 Query .newAggregationQueryBuilder ()
181170 .over (baseQuery )
182171 .addAggregations (avg ("marks" ).as ("avg_marks" ))
183- .setNamespace (OPTIONS .getNamespace ())
184172 .build ();
185173
186174 // avg of 2 entities
@@ -199,11 +187,7 @@ public void testAvgAggregationWithAutoGeneratedAlias() {
199187
200188 EntityQuery baseQuery = Query .newEntityQueryBuilder ().setKind (KIND ).build ();
201189 AggregationQuery aggregationQuery =
202- Query .newAggregationQueryBuilder ()
203- .over (baseQuery )
204- .addAggregations (avg ("marks" ))
205- .setNamespace (OPTIONS .getNamespace ())
206- .build ();
190+ Query .newAggregationQueryBuilder ().over (baseQuery ).addAggregations (avg ("marks" )).build ();
207191
208192 // avg of 2 entities
209193 assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getDouble ("property_1" ))
@@ -223,11 +207,7 @@ public void testAvgAggregationInGqlQuery() {
223207 Query .newGqlQueryBuilder ("AGGREGATE AVG(marks) AS avg_marks OVER (SELECT * FROM Marks)" )
224208 .build ();
225209
226- AggregationQuery aggregationQuery =
227- Query .newAggregationQueryBuilder ()
228- .over (gqlQuery )
229- .setNamespace (OPTIONS .getNamespace ())
230- .build ();
210+ AggregationQuery aggregationQuery = Query .newAggregationQueryBuilder ().over (gqlQuery ).build ();
231211
232212 // avg of 2 entities
233213 assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getDouble ("avg_marks" ))
@@ -249,7 +229,6 @@ public void testSumAndAvgAggregationTogether() {
249229 .over (baseQuery )
250230 .addAggregations (sum ("marks" ).as ("total_marks" ))
251231 .addAggregations (avg ("marks" ).as ("avg_marks" ))
252- .setNamespace (OPTIONS .getNamespace ())
253232 .build ();
254233
255234 // sum of 2 entities
@@ -271,7 +250,6 @@ public void testTransactionShouldReturnAConsistentSnapshot() {
271250 .addAggregation (count ().as ("count" ))
272251 .addAggregations (sum ("marks" ).as ("total_marks" ))
273252 .addAggregations (avg ("marks" ).as ("avg_marks" ))
274- .setNamespace (OPTIONS .getNamespace ())
275253 .build ();
276254
277255 // original entity count is 2
@@ -332,7 +310,6 @@ public void testReadOnlyTransactionShouldNotLockTheDocuments()
332310 .addAggregation (count ().as ("count" ))
333311 .addAggregations (sum ("marks" ).as ("total_marks" ))
334312 .addAggregations (avg ("marks" ).as ("avg_marks" ))
335- .setNamespace (OPTIONS .getNamespace ())
336313 .build ();
337314
338315 TransactionOptions transactionOptions =
0 commit comments