1 files changed, 94 insertions, 0 deletions
diff --git a/python/mozperftest/mozperftest/system/pingserver.py b/python/mozperftest/mozperftest/system/pingserver.py
new file mode 100644
index 0000000000..4ae6b9a113
--- /dev/null
+++ b/python/mozperftest/mozperftest/system/pingserver.py
@@ -0,0 +1,94 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+import json
+import os
+import socketserver
+import threading
+import time
+from pathlib import Path
+
+from mozlog import get_proxy_logger
+
+from mozperftest.layers import Layer
+from mozperftest.utils import install_package
+
+LOG = get_proxy_logger(component="proxy")
+HERE = os.path.dirname(__file__)
+
+
+class PingServer(Layer):
+    """Runs the edgeping layer"""
+
+    name = "pingserver"
+    activated = False
+
+    arguments = {}
+
+    def setup(self):
+        # Install edgeping and requests
+        deps = ["edgeping==0.1", "requests==2.9.1"]
+        for dep in deps:
+            install_package(self.mach_cmd.virtualenv_manager, dep)
+
+    def _wait_for_server(self, endpoint):
+        import requests
+
+        start = time.monotonic()
+        while True:
+            try:
+                requests.get(endpoint, timeout=0.1)
+                return
+            except Exception:
+                # we want to wait at most 5sec.
+                if time.monotonic() - start > 5.0:
+                    raise
+                time.sleep(0.01)
+
+    def run(self, metadata):
+        from edgeping.server import PingHandling
+
+        self.verbose = self.get_arg("verbose")
+        self.metadata = metadata
+        self.debug("Starting the Edgeping server")
+        self.httpd = socketserver.TCPServer(("localhost", 0), PingHandling)
+        self.server_thread = threading.Thread(target=self.httpd.serve_forever)
+        # the chosen socket gets picked in the constructor so we can grab it here
+        address = self.httpd.server_address
+        self.endpoint = f"http://{address[0]}:{address[1]}"
+        self.server_thread.start()
+        self._wait_for_server(self.endpoint + "/status")
+
+        self.debug(f"Edgeping coserver running at {self.endpoint}")
+        prefs = {
+            "toolkit.telemetry.server": self.endpoint,
+            "telemetry.fog.test.localhost_port": address[1],
+            "datareporting.healthreport.uploadEnabled": True,
+            "datareporting.policy.dataSubmissionEnabled": True,
+            "toolkit.telemetry.enabled": True,
+            "toolkit.telemetry.unified": True,
+            "toolkit.telemetry.shutdownPingSender.enabled": True,
+            "datareporting.policy.dataSubmissionPolicyBypassNotification": True,
+            "toolkit.telemetry.send.overrideOfficialCheck": True,
+        }
+        if self.verbose:
+            prefs["toolkit.telemetry.log.level"] = "Trace"
+            prefs["toolkit.telemetry.log.dump"] = True
+
+        browser_prefs = metadata.get_options("browser_prefs")
+        browser_prefs.update(prefs)
+        return metadata
+
+    def teardown(self):
+        import requests
+
+        self.info("Grabbing the pings")
+        pings = requests.get(f"{self.endpoint}/pings").json()
+        output = Path(self.get_arg("output"), "telemetry.json")
+        self.info(f"Writing in {output}")
+        with output.open("w") as f:
+            f.write(json.dumps(pings))
+
+        self.debug("Stopping the Edgeping coserver")
+        self.httpd.shutdown()
+        self.server_thread.join()