From b6566e395b2e03b3337b0f7ffacc2ed07b76e54f Mon Sep 17 00:00:00 2001
From: Alberto Soutullo <alberto.soutullo639@gmail.com>
Date: Mon, 24 Jul 2023 11:01:30 +0200
Subject: [PATCH] Added range to prometheus dumping data.

---
 wls-module/src/utils/prometheus.py | 5 +++--
 wls-module/src/wls.py              | 8 +++++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/wls-module/src/utils/prometheus.py b/wls-module/src/utils/prometheus.py
index e30ddc6..4703bbd 100644
--- a/wls-module/src/utils/prometheus.py
+++ b/wls-module/src/utils/prometheus.py
@@ -16,7 +16,7 @@ def connect_to_prometheus(ip, port):
     return prometheus
 
 
-def dump_prometheus(config, prometheus_ip, prometheus_port):
+def dump_prometheus(config, prometheus_ip, prometheus_port, start_time, finish_time):
     to_query = config["plotting"]["by_node"]
 
     to_query = "|".join(to_query)
@@ -28,7 +28,8 @@ def dump_prometheus(config, prometheus_ip, prometheus_port):
     query = f"{{__name__=~\"{to_query}\"}}"
     print(query)
 
-    metrics = prometheus_connection.custom_query(query)
+    metrics = prometheus_connection.custom_query_range(query, start_time=start_time,
+                                                       end_time=finish_time, step="1s")
 
     with open("/wls/prometheus_data.json", "w") as out_file:
         json.dump(metrics, out_file)
diff --git a/wls-module/src/wls.py b/wls-module/src/wls.py
index 39a30f7..3a6795f 100644
--- a/wls-module/src/wls.py
+++ b/wls-module/src/wls.py
@@ -7,6 +7,7 @@ import time
 import tomllib
 import asyncio
 import os
+from datetime import datetime
 
 # Project Imports
 from src.utils import wls_logger
@@ -197,8 +198,12 @@ async def main():
     t1 = time.time()
     wls_logger.G_LOGGER.info(f'Got the signal to start: took {t1-t0} secs')
 
+    injection_start_time = datetime.now()
+
     msgs_dict = await start_traffic_injection_async(wls_config, random_emitters)
 
+    injection_finish_time = datetime.now()
+
     files.save_messages_to_json(msgs_dict)
 
     # Delete de signal file just in case
@@ -206,7 +211,8 @@ async def main():
         os.remove('/wls/start.signal')
 
     if prometheus_port is not None:
-        prometheus.dump_prometheus(config, prometheus_ip, prometheus_port)
+        prometheus.dump_prometheus(config, prometheus_ip, prometheus_port, injection_start_time,
+                                   injection_finish_time)
 
 
 if __name__ == "__main__":