fix: Limit to one backend

platform_plugin_aspects/management/commands/monitor_load_test_tracking.py

@@ -24,7 +24,7 @@
 
 
 class Monitor:
-    def __init__(self, sleep_time: float):
+    def __init__(self, sleep_time: float, backend: str):
         self.run_id = str(uuid.uuid4())[:6]
         self.ch_url = settings.EVENT_SINK_CLICKHOUSE_BACKEND_CONFIG["url"]
         self.ch_auth = ClickHouseAuth(
@@ -39,13 +39,17 @@ def __init__(self, sleep_time: float):
             "timeout_secs"
         ]
         self.sleep_time = sleep_time
+        self.backend = backend
 
     def run(self) -> None:
-        collect_redis_bus = True
-        collect_celery = True
-        collect_kafka_bus = False
+        collect_redis_bus = True if self.backend == "redis_bus" else False
+        collect_celery = True if self.backend == "celery" else False
+        collect_kafka_bus = True if self.backend == "kafka_bus" else False
+        collect_vector = True if self.backend == "vector" else False
 
         while True:
+            print(f"----------- {datetime.datetime.now()} --------")
+
             # TODO: Need to pass in which stats to collect, or infer them from config
             current_stats = {"clickhouse": self.get_clickhouse_stats()}
             if collect_redis_bus:
@@ -54,6 +58,8 @@ def run(self) -> None:
                 current_stats["celery"] = self.get_celery_stats()
             # if collect_kafka_bus:
                 # current_stats["kafka_bus"] = self.get_kafka_stats()
+            # if collect_vector:
+                # current_stats["vector"] = self.get_vector_stats()
 
             self.store_stats(current_stats)
 
@@ -84,7 +90,6 @@ def store_stats(self, current_stats: dict) -> None:
         )
 
         response.raise_for_status()
-        print(datetime.datetime.now())
 
     def get_clickhouse_stats(self):
         select = f"""
@@ -192,13 +197,19 @@ def add_arguments(self, parser: Any) -> None:
             default=10,
             help="Fractional number of seconds to sleep between gathering data.",
         )
+        parser.add_argument(
+            "--backend",
+            choices=["redis_bus", "kafka_bus", "celery", "vector"],
+            default="celery",
+            help="Backend used to send events to ClickHouse",
+        )
 
     def handle(self, *args, **options):
         """
         Creates users and triggers events for them as configured above.
         """
         start = datetime.datetime.now()
-        monitor = Monitor(options["sleep_time"])
+        monitor = Monitor(options["sleep_time"], options["backend"])
 
         try:
             monitor.run()