Merge pull request #48 from vinod-kanigicherla/fix-netflex-pipelines

Add handlers for Ookla and MLab; fixed measurement_data formatting

Author: maybe-hello-world

pipelines/measurements/netflex_pipeline.py

@@ -3,6 +3,8 @@
 from netunicorn.library.tasks.measurements.ndt import NDT7SpeedTest
 from netunicorn.library.tasks.data_transfer.send_data import SendData
 from netunicorn.library.tasks.data_transfer.fetch_data import FetchData
+from netunicorn.library.tasks.handlers.ookla import ookla_handler
+from netunicorn.library.tasks.handlers.mlab import mlab_handler
 import os 
 
 def netflex_ookla_full_loop_pipeline() -> Pipeline:
@@ -14,7 +16,7 @@ def netflex_ookla_full_loop_pipeline() -> Pipeline:
     pipeline = (
         Pipeline()
         .then(OoklaSpeedtest(name="Ookla CLI Speedtest"))
-        .then(SendData(task_descriptors=[SendData.TaskDescriptor("Ookla CLI Speedtest", "ookla-speedtest", None)], endpoint=os.getenv('RAG_ENDPOINT'), name="send"))
+        .then(SendData(task_descriptors=[SendData.TaskDescriptor("Ookla CLI Speedtest", "ookla-speedtest", ookla_handler)], endpoint=os.getenv('RAG_ENDPOINT'), name="send"))
         .then(FetchData(send_data_task="send", endpoint=os.getenv("RAG_ENDPOINT"), name="fetch"))
     )
     return pipeline
@@ -28,7 +30,7 @@ def netflex_mlab_full_loop_pipeline() -> Pipeline:
     pipeline = (
         Pipeline()
         .then(NDT7SpeedTest(name="test", flags=["-format", "json"]))
-        .then(SendData(task_descriptors=[SendData.TaskDescriptor("test", "mlab-speedtest", None)], endpoint=os.getenv('RAG_ENDPOINT'), name="send"))
+        .then(SendData(task_descriptors=[SendData.TaskDescriptor("test", "mlab-speedtest", mlab_handler)], endpoint=os.getenv('RAG_ENDPOINT'), name="send"))
         .then(FetchData(send_data_task="send", endpoint=os.getenv("RAG_ENDPOINT"), name="fetch"))
     )
     return pipeline

tasks/data_transfer/send_data.py

@@ -18,14 +18,16 @@ def __init__(
         self, 
         task_descriptors: list[TaskDescriptor], 
         endpoint: str, 
-        allow_failure = False, 
+        allow_failure = False,
+	payload_handler: Optional[Callable[[Any], Any]] = None,
         *args, 
         **kwargs
         ):
 
         self.task_descriptors = task_descriptors
         self.endpoint = endpoint
         self.allow_failure = allow_failure
+        self.payload_handler = payload_handler
         super().__init__(*args, **kwargs)
 
     def get_geolocation_from_ip(self, ip: str): 
@@ -61,21 +63,31 @@ def run(self):
                 if isinstance(raw_task_result, Failure):
                     continue
 
-                task_results = [result.unwrap() for result in raw_task_result]
+                raw_results = [result.unwrap() for result in raw_task_result]
 
                 if task_descriptor.handler:
-                    task_results = task_descriptor.handler(task_descriptor)
-                
-                execution_results.append({"task_type": task_descriptor.datatype, "task_results": task_results})
+                    task_results = task_descriptor.handler(task_descriptor.datatype, raw_results)
+                else: # no handler
+                    task_results = {
+                    "task_type": task_descriptor.datatype,
+                    "task_results": raw_results
+                    }
+                execution_results.append(task_results)
 
             (location_str, *_) = self.get_geolocation_from_ip("8.8.8.8")
 
-            response = requests.post(self.endpoint, json=
-	       {
-                   "client_location": location_str,
-                   "measurement_type": self.task_descriptors[0].datatype,
-                   "measurement_data": execution_results
-               })
+            measurement_data = {"execution_results": execution_results}
+
+            payload = {
+                    "client_location": location_str,
+            "measurement_type": self.task_descriptors[0].datatype,
+            "measurement_data": measurement_data
+                }
+
+            if self.payload_handler:
+                    payload = self.payload_handler(payload)
+ 
+            response = requests.post(self.endpoint, json=payload)
 
             if response.status_code == 200:
                 return Success(response.json())

tasks/handlers/__init__.py

@@ -0,0 +1,2 @@
+def mlab_handler(task_type: str, results: list[dict]) -> dict:
+    return {"task_type": task_type, "task_results": {"data": results}}

tasks/handlers/mlab.py

@@ -0,0 +1,4 @@
+def ookla_handler(task_type: str, results: list[dict]) -> dict:
+    for r in results:
+        r.pop("interface", None) 
+    return {"task_type": task_type, "task_results": results}