fixed Failure errors, made output_file optional, and added new curl timers

Author: evaniacheng

tasks/measurements/alexa/alexa.py

@@ -1,7 +1,7 @@
 from typing import Dict, Union
 from netunicorn.base import Task, Failure
-from pprint import pprint
 import subprocess
+import pprint
 from ping3 import ping
 import csv
 import json
@@ -13,7 +13,7 @@ class AlexaWebsitesTask(Task):
         "pip install ping3"
     ]
 
-    def __init__(self, domain: str = None, filepath: str = None, output_path: str = None, top_k: int = 100, *args, **kwargs):
+    def __init__(self, domain: str = None, filepath: str = "alexa_websites.csv", output_path: str = None, top_k: int = 100, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.domain = domain
         self.filepath = filepath
@@ -31,7 +31,7 @@ def measure_ping(self) -> Union[Dict[str, float], Failure]:
         try:
             ping_value = ping(self.domain)
             if ping_value is None:
-                return Failure({"Ping returned None."}) 
+                return Failure("Ping returned None.") 
             return {"value": ping_value * 1000, "unit": "ms"}
         except Exception as e:
             return Failure(f"Ping failed: {e}")
@@ -46,20 +46,33 @@ def measure_dns_time(self) -> Union[Dict[str, float], Failure]:
         except Exception as e:
             return Failure(f"DNS resolution failed: {e}")
 
-    def measure_ttfb(self) -> Union[Dict[str, float], Failure]:
+    def measure_timing(self) -> Union[Dict[str, Dict[str,float]], Failure]:
         try:
             result = subprocess.run([
                 "curl",
                 "-o", "/dev/null",
                 "-s",
-                "-w", "%{time_starttransfer}",
+                "-w", 
+                (
+                    "time_appconnect: %{time_appconnect}\n"
+                    "time_connect: %{time_connect}\n"
+                    "time_namelookup: %{time_namelookup}\n"
+                    "time_pretransfer: %{time_pretransfer}\n"
+                    "time_redirect: %{time_redirect}\n"
+                    "time_starttransfer: %{time_starttransfer}\n"
+                    "time_total: %{time_total}\n"
+                ),
                 "-H", "Cache-Control: no-cache",
-                f"http://{self.domain}",
+                f"https://{self.domain}",
             ], capture_output=True, text=True, check=True)
-            ttfb = float(result.stdout.strip()) * 1000  # s to ms
-            return {"value": ttfb, "unit": "ms"}
+            metrics = {
+                key.strip(): {"value": float(value.strip()) * 1000, "unit": "ms"}
+                for line in result.stdout.splitlines()
+                for key, value in [line.split(": ", 1)]
+            }
+            return metrics
         except Exception as e:
-            return Failure(f"TTFB measurement failed: {e}")
+            return Failure(f"Network Timing measurement failed: {e}")
 
     @staticmethod
     def load_websites(filepath: str, top_k: int) -> list:
@@ -78,16 +91,14 @@ def run(self) -> Union[Dict[str, Dict], Failure]:
         if self.domain:
             # Run for a single domain
             return {
-                "domain": self.domain,
                 "traceroute": self.get_traceroute(),
                 "ping_time": self.measure_ping(),
                 "dns_time": self.measure_dns_time(),
-                "ttfb_time": self.measure_ttfb(),
+                "measure_timing": self.measure_timing(),
             }
-
-        elif self.filepath and self.output_path:
-            # Run for all websites in the file
-            websites = AlexaWebsitesTask.load_websites(self.filepath, self.top_k)
+        else:
+            # Run for all websites in a file
+            websites = self.load_websites(self.filepath, self.top_k)
             print(f"Loaded {len(websites)} websites.")
 
             results = {}
@@ -99,12 +110,15 @@ def run(self) -> Union[Dict[str, Dict], Failure]:
                 except Exception as e:
                     results[website] = Failure(f"Failed to process {website}: {e}")
 
-            # Save results to a JSON file
-            with open(self.output_path, "w") as f:
-                json.dump(results, f, indent=4)
-
-            print(f"Results saved to {self.output_path}")
+            # Save results to a JSON file if output_path is provided
+            if self.output_path:
+                print(f"Saving results to {self.output_path}")
+                try: 
+                    with open(self.output_path, "w") as f:
+                        json.dump(results, f, indent=4)
+                except Exception as e:
+                    return Failure(f"Failed to write results to file: {e}") 
+            else:
+                pprint.pp(results)
+                
             return results
-
-        else:
-            raise Failure("Either a domain or both a filepath and output_path must be provided.")