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 ;
2930import  com .google .cloud .datastore .Entity ;
3031import  com .google .cloud .datastore .EntityQuery ;
3132import  com .google .cloud .datastore .GqlQuery ;
3233import  com .google .cloud .datastore .Key ;
3334import  com .google .cloud .datastore .Query ;
3435import  com .google .cloud .datastore .QueryResults ;
3536import  com .google .cloud .datastore .Transaction ;
37+ import  com .google .cloud .datastore .testing .RemoteDatastoreHelper ;
3638import  com .google .common .collect .ImmutableList ;
3739import  com .google .datastore .v1 .TransactionOptions ;
3840import  com .google .datastore .v1 .TransactionOptions .ReadOnly ;
4143import  java .util .concurrent .Executors ;
4244import  java .util .concurrent .Future ;
4345import  org .junit .After ;
44- import  org .junit .BeforeClass ;
45- import  org .junit .ClassRule ;
46+ import  org .junit .AfterClass ;
4647import  org .junit .Test ;
4748
4849// TODO(jainsahab) Move all the aggregation related tests from ITDatastoreTest to this file 
4950public  class  ITDatastoreAggregationsTest  {
5051
51-   @ ClassRule   public   static  MultiDbRule   multiDbRule  =  new   MultiDbRule ();
52- 
53-   private  static  Datastore  DATASTORE ;
52+   private   static  final   RemoteDatastoreHelper   HELPER  =  RemoteDatastoreHelper . create ();
53+    private   static   final   DatastoreOptions   OPTIONS  =  HELPER . getOptions (); 
54+   private  static  final   Datastore  DATASTORE  =  OPTIONS . getService () ;
5455
5556  private  static  final  String  KIND  = "Marks" ;
5657
57-   @ BeforeClass 
58-   public  static  void  beforeClass () throws  Exception  {
59-     DATASTORE  = multiDbRule .getDatastore ();
60-   }
61- 
6258  @ After 
6359  public  void  tearDown () {
6460    EntityQuery  allEntitiesQuery  = Query .newEntityQueryBuilder ().build ();
@@ -68,6 +64,11 @@ public void tearDown() {
6864    DATASTORE .delete (keysToDelete );
6965  }
7066
67+   @ AfterClass 
68+   public  static  void  afterClass () throws  Exception  {
69+     DATASTORE .close ();
70+   }
71+ 
7172  Key  key1  = DATASTORE .newKeyFactory ().setKind (KIND ).newKey (1 );
7273  Key  key2  = DATASTORE .newKeyFactory ().setKind (KIND ).newKey (2 );
7374  Key  key3  = DATASTORE .newKeyFactory ().setKind (KIND ).newKey (3 );
@@ -88,6 +89,7 @@ public void testSumAggregation() {
8889        Query .newAggregationQueryBuilder ()
8990            .over (baseQuery )
9091            .addAggregations (sum ("marks" ).as ("total_marks" ))
92+             .setNamespace (OPTIONS .getNamespace ())
9193            .build ();
9294
9395    // sum of 2 entities 
@@ -106,7 +108,11 @@ public void testSumAggregationWithAutoGeneratedAlias() {
106108
107109    EntityQuery  baseQuery  = Query .newEntityQueryBuilder ().setKind (KIND ).build ();
108110    AggregationQuery  aggregationQuery  =
109-         Query .newAggregationQueryBuilder ().over (baseQuery ).addAggregations (sum ("marks" )).build ();
111+         Query .newAggregationQueryBuilder ()
112+             .over (baseQuery )
113+             .addAggregations (sum ("marks" ))
114+             .setNamespace (OPTIONS .getNamespace ())
115+             .build ();
110116
111117    // sum of 2 entities 
112118    assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getLong ("property_1" ))
@@ -127,7 +133,11 @@ public void testSumAggregationInGqlQuery() {
127133                "AGGREGATE SUM(marks) AS total_marks OVER (SELECT * FROM Marks)" )
128134            .build ();
129135
130-     AggregationQuery  aggregationQuery  = Query .newAggregationQueryBuilder ().over (gqlQuery ).build ();
136+     AggregationQuery  aggregationQuery  =
137+         Query .newAggregationQueryBuilder ()
138+             .over (gqlQuery )
139+             .setNamespace (OPTIONS .getNamespace ())
140+             .build ();
131141
132142    // sum of 2 entities 
133143    assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getLong ("total_marks" ))
@@ -148,6 +158,7 @@ public void testSumAggregationWithResultOfDoubleType() {
148158        Query .newAggregationQueryBuilder ()
149159            .over (baseQuery )
150160            .addAggregations (sum ("cgpa" ).as ("total_cgpa" ))
161+             .setNamespace (OPTIONS .getNamespace ())
151162            .build ();
152163
153164    // sum of 2 entities 
@@ -169,6 +180,7 @@ public void testAvgAggregation() {
169180        Query .newAggregationQueryBuilder ()
170181            .over (baseQuery )
171182            .addAggregations (avg ("marks" ).as ("avg_marks" ))
183+             .setNamespace (OPTIONS .getNamespace ())
172184            .build ();
173185
174186    // avg of 2 entities 
@@ -187,7 +199,11 @@ public void testAvgAggregationWithAutoGeneratedAlias() {
187199
188200    EntityQuery  baseQuery  = Query .newEntityQueryBuilder ().setKind (KIND ).build ();
189201    AggregationQuery  aggregationQuery  =
190-         Query .newAggregationQueryBuilder ().over (baseQuery ).addAggregations (avg ("marks" )).build ();
202+         Query .newAggregationQueryBuilder ()
203+             .over (baseQuery )
204+             .addAggregations (avg ("marks" ))
205+             .setNamespace (OPTIONS .getNamespace ())
206+             .build ();
191207
192208    // avg of 2 entities 
193209    assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getDouble ("property_1" ))
@@ -207,7 +223,11 @@ public void testAvgAggregationInGqlQuery() {
207223        Query .newGqlQueryBuilder ("AGGREGATE AVG(marks) AS avg_marks OVER (SELECT * FROM Marks)" )
208224            .build ();
209225
210-     AggregationQuery  aggregationQuery  = Query .newAggregationQueryBuilder ().over (gqlQuery ).build ();
226+     AggregationQuery  aggregationQuery  =
227+         Query .newAggregationQueryBuilder ()
228+             .over (gqlQuery )
229+             .setNamespace (OPTIONS .getNamespace ())
230+             .build ();
211231
212232    // avg of 2 entities 
213233    assertThat (getOnlyElement (DATASTORE .runAggregation (aggregationQuery )).getDouble ("avg_marks" ))
@@ -229,6 +249,7 @@ public void testSumAndAvgAggregationTogether() {
229249            .over (baseQuery )
230250            .addAggregations (sum ("marks" ).as ("total_marks" ))
231251            .addAggregations (avg ("marks" ).as ("avg_marks" ))
252+             .setNamespace (OPTIONS .getNamespace ())
232253            .build ();
233254
234255    // sum of 2 entities 
@@ -250,6 +271,7 @@ public void testTransactionShouldReturnAConsistentSnapshot() {
250271            .addAggregation (count ().as ("count" ))
251272            .addAggregations (sum ("marks" ).as ("total_marks" ))
252273            .addAggregations (avg ("marks" ).as ("avg_marks" ))
274+             .setNamespace (OPTIONS .getNamespace ())
253275            .build ();
254276
255277    // original entity count is 2 
@@ -310,6 +332,7 @@ public void testReadOnlyTransactionShouldNotLockTheDocuments()
310332            .addAggregation (count ().as ("count" ))
311333            .addAggregations (sum ("marks" ).as ("total_marks" ))
312334            .addAggregations (avg ("marks" ).as ("avg_marks" ))
335+             .setNamespace (OPTIONS .getNamespace ())
313336            .build ();
314337
315338    TransactionOptions  transactionOptions  =
0 commit comments