refac: refactor strategy selection logic

wazuh · Jun 25, 2024 · 0312c7e · 0312c7e
1 parent e36d0a1
commit 0312c7e
Showing 1 changed file with 13 additions and 11 deletions.
diff --git a/deps/wazuh_testing/wazuh_testing/scripts/data_visualizations.py b/deps/wazuh_testing/wazuh_testing/scripts/data_visualizations.py
@@ -18,6 +18,7 @@
                      'indexer-vulnerabilities']
 strategy_plot_by_target = {
     'binary': BinaryDatavisualizer,
+    'daemon-statistics': DaemonStatisticsVisualizer,
     'cluster': ClusterStatisticsVisualizer,
     'logcollector': LogcollectorStatisticsVisualizer,
     'indexer-alerts': IndexerAlerts,
@@ -61,19 +62,20 @@ def main():
     target = options.visualization_target
     validate_arguments(options)
 
+    visualization_options = {
+        'dataframes': options.csv_list,
+        'store_path': options.destination,
+        'base_name': options.name
+    }
+
+    strategy = target
     if target in ['analysis', 'remote', 'wazuhdb']:
-        dv = DaemonStatisticsVisualizer(options.csv_list, daemon=target,
-                                        store_path=options.destination,
-                                        base_name=options.name)
+        visualization_options['daemon'] = target
+        strategy = 'daemon-statistics'
     elif target == 'binary':
-        dv = BinaryDatavisualizer(options.csv_list,
-                                    store_path=options.destination,
-                                    base_name=options.name,
-                                    unify_child_daemon_metrics=options.unify)
-    else:
-        dv = strategy_plot_by_target[target](options.csv_list,
-                                        store_path=options.destination,
-                                        base_name=options.name)
+        visualization_options['unify_child_daemon_metrics'] = options.unify
+
+    dv = strategy_plot_by_target[strategy](**visualization_options)
 
     dv.plot()