fix: Using AuthorizedHttpClient for updating rules during tests (#590)

Author: hiranya911

src/test/java/com/google/firebase/database/integration/FirebaseDatabaseAuthTestIT.java

@@ -16,7 +16,6 @@
 
 package com.google.firebase.database.integration;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 import com.google.api.core.ApiFutureCallback;
@@ -33,8 +32,6 @@
 import com.google.firebase.database.TestHelpers;
 import com.google.firebase.database.ValueEventListener;
 import com.google.firebase.testing.IntegrationTestUtils;
-import com.google.firebase.testing.IntegrationTestUtils.AppHttpClient;
-import com.google.firebase.testing.IntegrationTestUtils.ResponseInfo;
 import com.google.firebase.testing.ServiceAccount;
 import com.google.firebase.testing.TestUtils;
 import java.io.IOException;
@@ -49,11 +46,11 @@
 import org.junit.Test;
 
 public class FirebaseDatabaseAuthTestIT {
-  
+
   private static FirebaseApp masterApp;
-  
+
   @BeforeClass
-  public static void setUpClass() throws IOException {    
+  public static void setUpClass() throws IOException {
     masterApp = IntegrationTestUtils.ensureDefaultApp();
     setDatabaseRules();
   }
@@ -74,7 +71,7 @@ public void testAuthWithValidCertificateCredential() throws InterruptedException
     assertWriteSucceeds(db.getReference());
     assertReadSucceeds(db.getReference());
   }
-  
+
   @Test
   public void testAuthWithInvalidCertificateCredential() throws InterruptedException, IOException {
     FirebaseOptions options =
@@ -87,7 +84,7 @@ public void testAuthWithInvalidCertificateCredential() throws InterruptedExcepti
     // TODO: Ideally, we would find a way to verify the correct log output.
     assertWriteTimeout(db.getReference());
   }
-  
+
   @Test
   public void testDatabaseAuthVariablesAuthorization() throws InterruptedException {
     Map<String, Object> authVariableOverrides = ImmutableMap.<String, Object>of(
@@ -110,7 +107,7 @@ public void testDatabaseAuthVariablesAuthorization() throws InterruptedException
     assertWriteSucceeds(testAuthOverridesDb.getReference("test-custom-field-only"));
     assertReadSucceeds(testAuthOverridesDb.getReference("test-custom-field-only"));
   }
-  
+
   @Test
   public void testDatabaseAuthVariablesNoAuthorization() throws InterruptedException {
     FirebaseOptions options = masterApp.getOptions().toBuilder()
@@ -128,25 +125,25 @@ public void testDatabaseAuthVariablesNoAuthorization() throws InterruptedExcepti
     assertReadFails(testAuthOverridesDb.getReference("test-uid-only"));
     assertWriteFails(testAuthOverridesDb.getReference("test-custom-field-only"));
     assertReadFails(testAuthOverridesDb.getReference("test-custom-field-only"));
-    assertWriteSucceeds(testAuthOverridesDb.getReference("test-noauth-only"));    
+    assertWriteSucceeds(testAuthOverridesDb.getReference("test-noauth-only"));
   }
-  
+
   private static void assertWriteSucceeds(DatabaseReference ref) throws InterruptedException {
     doWrite(ref, /*shouldSucceed=*/ true, /*shouldTimeout=*/ false);
   }
-  
+
   private static void assertWriteFails(DatabaseReference ref) throws InterruptedException {
     doWrite(ref, /*shouldSucceed=*/ false, /*shouldTimeout=*/ false);
   }
-  
+
   private static void assertWriteTimeout(DatabaseReference ref) throws InterruptedException {
     doWrite(ref, /*shouldSucceed=*/ false, /*shouldTimeout=*/ true);
   }
-  
+
   private static void assertReadSucceeds(DatabaseReference ref) throws InterruptedException {
     doRead(ref, /*shouldSucceed=*/ true, /*shouldTimeout=*/ false);
   }
-  
+
   private static void assertReadFails(DatabaseReference ref) throws InterruptedException {
     doRead(ref, /*shouldSucceed=*/ false, /*shouldTimeout=*/ false);
   }
@@ -180,7 +177,7 @@ public void onSuccess(Void result) {
       assertTrue("Write successful (expected to fail).", !success.get());
     }
   }
-  
+
   private static void doRead(
       DatabaseReference ref, final boolean shouldSucceed, final boolean shouldTimeout)
       throws InterruptedException {
@@ -211,9 +208,9 @@ public void onCancelled(DatabaseError databaseError) {
       assertTrue("Read successful (expected to fail).", !success.get());
     }
   }
-  
+
   private static void setDatabaseRules() throws IOException {
-    // TODO: Use more than uid in rule Set rules so the only allowed operation is writing to 
+    // TODO: Use more than uid in rule Set rules so the only allowed operation is writing to
     // /test-uid-only by user with uid 'test'.
     String rules =
         "{\n"
@@ -232,8 +229,7 @@ private static void setDatabaseRules() throws IOException {
             + "  }\n"
             + "}";
 
-    AppHttpClient client = new AppHttpClient();
-    ResponseInfo info = client.put("/.settings/rules.json", rules);
-    assertEquals(200, info.getStatus());
+    RulesClient client = new RulesClient();
+    client.updateRules(rules);
   }
 }

src/test/java/com/google/firebase/database/integration/OrderByTestIT.java

@@ -87,9 +87,8 @@ private static String formatRules(DatabaseReference ref, String rules) {
   }
 
   private static void uploadRules(FirebaseApp app, String rules) throws IOException {
-    IntegrationTestUtils.AppHttpClient client = new IntegrationTestUtils.AppHttpClient(app);
-    IntegrationTestUtils.ResponseInfo response = client.put("/.settings/rules.json", rules);
-    assertEquals(200, response.getStatus());
+    RulesClient client = new RulesClient(app);
+    client.updateRules(rules);
   }
 
   @Test

src/test/java/com/google/firebase/database/integration/RulesClient.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2021 Google Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.firebase.database.integration;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.api.client.http.ByteArrayContent;
+import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpRequest;
+import com.google.api.client.http.HttpRequestFactory;
+import com.google.api.client.http.HttpResponse;
+import com.google.common.io.ByteStreams;
+import com.google.firebase.FirebaseApp;
+import com.google.firebase.internal.ApiClientUtils;
+import java.io.IOException;
+import java.io.InputStream;
+
+/**
+ * A simple HTTP client wrapper for updating RTDB security rules.
+ */
+final class RulesClient {
+
+  private final String databaseUrl;
+  private final HttpRequestFactory requestFactory;
+
+  RulesClient() {
+    this(FirebaseApp.getInstance());
+  }
+
+  RulesClient(FirebaseApp app) {
+    this.databaseUrl = checkNotNull(app).getOptions().getDatabaseUrl();
+    this.requestFactory = ApiClientUtils.newAuthorizedRequestFactory(
+      app, /*retryConfig*/ null);
+  }
+
+  public void updateRules(String json) throws IOException {
+    String url = this.databaseUrl + "/.settings/rules.json";
+    HttpRequest request = requestFactory.buildPutRequest(
+        new GenericUrl(url),
+        ByteArrayContent.fromString("application/json", json));
+    HttpResponse response = request.execute();
+    try {
+      InputStream in = response.getContent();
+      if (in != null) {
+        ByteStreams.exhaust(in);
+      }
+    } finally {
+      ApiClientUtils.disconnectQuietly(response);
+    }
+  }
+}

src/test/java/com/google/firebase/database/integration/RulesTestIT.java

@@ -66,7 +66,7 @@ public class RulesTestIT {
       MapBuilder.of("rules", MapBuilder.of(".read", "auth != null", ".write", "auth != null"));
 
   private static final Map<String, Object> testRules;
-  
+
   static {
     testRules = new MapBuilder()
         .put("read_only", MapBuilder.of(".read", true))
@@ -137,9 +137,8 @@ public void checkAndCleanupApp() {
   }
 
   private static void uploadRules(String rules) throws IOException {
-    IntegrationTestUtils.AppHttpClient client = new IntegrationTestUtils.AppHttpClient(masterApp);
-    IntegrationTestUtils.ResponseInfo response = client.put("/.settings/rules.json", rules);
-    assertEquals(200, response.getStatus());
+    RulesClient client = new RulesClient(masterApp);
+    client.updateRules(rules);
   }
 
   @Test

src/test/java/com/google/firebase/testing/IntegrationTestUtils.java

@@ -16,21 +16,12 @@
 
 package com.google.firebase.testing;
 
-import static com.google.common.base.Preconditions.checkNotNull;
-
-import com.google.api.client.http.ByteArrayContent;
-import com.google.api.client.http.GenericUrl;
-import com.google.api.client.http.HttpRequest;
-import com.google.api.client.http.HttpRequestFactory;
-import com.google.api.client.http.HttpResponse;
 import com.google.api.client.json.GenericJson;
 import com.google.cloud.firestore.FirestoreOptions;
 import com.google.common.collect.ImmutableList;
-import com.google.common.io.ByteStreams;
 import com.google.common.io.CharStreams;
 import com.google.firebase.FirebaseApp;
 import com.google.firebase.FirebaseOptions;
-import com.google.firebase.TestOnlyImplFirebaseTrampolines;
 import com.google.firebase.database.DatabaseReference;
 import com.google.firebase.database.FirebaseDatabase;
 import com.google.firebase.internal.ApiClientUtils;
@@ -148,65 +139,4 @@ public static List<DatabaseReference> getRandomNode(FirebaseApp app, int count)
     }
     return builder.build();
   }
-
-  public static class AppHttpClient {
-
-    private final FirebaseApp app;
-    private final FirebaseOptions options;
-    private final HttpRequestFactory requestFactory;
-
-    public AppHttpClient() {
-      this(FirebaseApp.getInstance());
-    }
-
-    public AppHttpClient(FirebaseApp app) {
-      this.app = checkNotNull(app);
-      this.options = app.getOptions();
-      this.requestFactory = this.options.getHttpTransport().createRequestFactory();
-    }
-
-    public ResponseInfo put(String path, String json) throws IOException {
-      String url = options.getDatabaseUrl() + path + "?access_token=" + getToken();
-      HttpRequest request = requestFactory.buildPutRequest(new GenericUrl(url),
-          ByteArrayContent.fromString("application/json", json));
-      HttpResponse response = null;
-      try {
-        response = request.execute();
-        return new ResponseInfo(response);
-      } finally {
-        if (response != null) {
-          response.disconnect();
-        }
-      }
-    }
-
-    private String getToken() {
-      // TODO: We should consider exposing getToken (or similar) publicly for the
-      // purpose of servers doing authenticated REST requests like this.
-      return TestOnlyImplFirebaseTrampolines.getToken(app, false);
-    }
-  }
-
-  public static class ResponseInfo {
-    private final int status;
-    private final byte[] payload;
-
-    private ResponseInfo(HttpResponse response) throws IOException {
-      this.status = response.getStatusCode();
-      InputStream in = response.getContent();
-      if (in != null) {
-        this.payload = ByteStreams.toByteArray(in);
-      } else {
-        this.payload = new byte[0];
-      }
-    }
-
-    public int getStatus() {
-      return status;
-    }
-
-    public byte[] getPayload() {
-      return payload;
-    }
-  }
 }