From 5d8823f1a6d3aab6291b6028f641f9788c20d385 Mon Sep 17 00:00:00 2001
From: Jan-Erik Rediger <jrediger@mozilla.com>
Date: Thu, 17 May 2018 01:02:00 +0200
Subject: [PATCH] Bug 1461652 - Mock gzip compression in tests to simulate
 large payloads. r=chutten

MozReview-Commit-ID: Gct9oVfPVou
---
 toolkit/components/telemetry/TelemetrySend.jsm    |  3 ++-
 toolkit/components/telemetry/tests/unit/head.js   | 15 +++++++++++++++
 .../telemetry/tests/unit/test_TelemetrySend.js    | 13 +++++++------
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/toolkit/components/telemetry/TelemetrySend.jsm b/toolkit/components/telemetry/TelemetrySend.jsm
index d0351cfc7853..c57f7adb0eb9 100644
--- a/toolkit/components/telemetry/TelemetrySend.jsm
+++ b/toolkit/components/telemetry/TelemetrySend.jsm
@@ -109,6 +109,7 @@ var Policy = {
   pingSubmissionTimeout: () => PING_SUBMIT_TIMEOUT_MS,
   setSchedulerTickTimeout: (callback, delayMs) => setTimeout(callback, delayMs),
   clearSchedulerTickTimeout: (id) => clearTimeout(id),
+  gzipCompressString: (data) => gzipCompressString(data),
 };
 
 /**
@@ -1214,7 +1215,7 @@ var TelemetrySendImpl = {
     let payloadStream = Cc["@mozilla.org/io/string-input-stream;1"]
                         .createInstance(Ci.nsIStringInputStream);
     startTime = Utils.monotonicNow();
-    payloadStream.data = gzipCompressString(utf8Payload);
+    payloadStream.data = Policy.gzipCompressString(utf8Payload);
 
     // Check the size and drop pings which are too big.
     const compressedPingSizeBytes = payloadStream.data.length;
diff --git a/toolkit/components/telemetry/tests/unit/head.js b/toolkit/components/telemetry/tests/unit/head.js
index 7b733270e36b..5fd535a0e7cd 100644
--- a/toolkit/components/telemetry/tests/unit/head.js
+++ b/toolkit/components/telemetry/tests/unit/head.js
@@ -275,6 +275,17 @@ function fakeCachedClientId(uuid) {
   module.Policy.getCachedClientID = () => uuid;
 }
 
+// Fake the gzip compression for the next ping to be sent out
+// and immediately reset to the original function.
+function fakeGzipCompressStringForNextPing(length) {
+  let send = ChromeUtils.import("resource://gre/modules/TelemetrySend.jsm", {});
+  let largePayload = generateString(length);
+  send.Policy.gzipCompressString = (data) => {
+    send.Policy.gzipCompressString = send.gzipCompressString;
+    return largePayload;
+  };
+}
+
 // Return a date that is |offset| ms in the future from |date|.
 function futureDate(date, offset) {
   return new Date(date.getTime() + offset);
@@ -301,6 +312,10 @@ function generateRandomString(length) {
   return string.substring(0, length);
 }
 
+function generateString(length) {
+  return new Array(length + 1).join("a");
+}
+
 // Short-hand for retrieving the histogram with that id.
 function getHistogram(histogramId) {
   return Telemetry.getHistogramById(histogramId);
diff --git a/toolkit/components/telemetry/tests/unit/test_TelemetrySend.js b/toolkit/components/telemetry/tests/unit/test_TelemetrySend.js
index ab0e9af55b38..04710afd29e6 100644
--- a/toolkit/components/telemetry/tests/unit/test_TelemetrySend.js
+++ b/toolkit/components/telemetry/tests/unit/test_TelemetrySend.js
@@ -334,9 +334,6 @@ add_task(async function test_discardBigPings() {
     h.clear();
   }
 
-  // Generate a 2MB string and create an oversized payload.
-  const OVERSIZED_PAYLOAD = {"data": generateRandomString(4 * 1024 * 1024)};
-
   // Submit a ping of a normal size and check that we don't count it in the histogram.
   await TelemetryController.submitExternalPing(TEST_PING_TYPE, { test: "test" });
   await TelemetrySend.testWaitOnOutgoingPings();
@@ -351,6 +348,9 @@ add_task(async function test_discardBigPings() {
 
   // Submit an oversized ping and check that it gets discarded.
   TelemetryHealthPing.testReset();
+  // Ensure next ping has a 2 MB gzipped payload.
+  fakeGzipCompressStringForNextPing(2 * 1024 * 1024);
+  const OVERSIZED_PAYLOAD = {"data": "empty on purpose - policy takes care of size"};
   await TelemetryController.submitExternalPing(TEST_PING_TYPE, OVERSIZED_PAYLOAD);
   await TelemetrySend.testWaitOnOutgoingPings();
   let ping = await PingServer.promiseNextPing();
@@ -376,12 +376,13 @@ add_task(async function test_largeButWithinLimit() {
   let histSuccess = Telemetry.getHistogramById("TELEMETRY_SUCCESS");
   histSuccess.clear();
 
-  // Generate a 900KB string and a large payload that is still within the 1MB limit.
-  const LARGE_PAYLOAD = {"data": generateRandomString(900 * 1024)};
+  // Next ping will have a 900KB gzip payload.
+  fakeGzipCompressStringForNextPing(900 * 1024);
+  const LARGE_PAYLOAD = {"data": "empty on purpose - policy takes care of size"};
 
   await TelemetryController.submitExternalPing(TEST_PING_TYPE, LARGE_PAYLOAD);
   await TelemetrySend.testWaitOnOutgoingPings();
-  await PingServer.promiseNextPing();
+  await PingServer.promiseNextRequest();
 
   Assert.deepEqual(histSuccess.snapshot().counts, [0, 1, 0], "Should have sent large ping.");
 });