From c16f33a111bf44e6b51f160007f9b4274408749e Mon Sep 17 00:00:00 2001
From: Mohse Morad <moshemorad12340@gmail.com>
Date: Thu, 7 Nov 2024 20:55:04 +0200
Subject: [PATCH] Enable exclude-severity only on csv format

---
 robusta_krr/formatters/csv.py          |  5 +++--
 robusta_krr/main.py                    |  3 +++
 tests/formatters/test_csv_formatter.py |  2 ++
 tests/test_krr.py                      |  2 +-
 tests/test_runner.py                   | 21 +++++++++++++++++++++
 5 files changed, 30 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_runner.py

diff --git a/robusta_krr/formatters/csv.py b/robusta_krr/formatters/csv.py
index 5e2ac95..d32f5f9 100644
--- a/robusta_krr/formatters/csv.py
+++ b/robusta_krr/formatters/csv.py
@@ -6,6 +6,7 @@
 
 from robusta_krr.core.abstract import formatters
 from robusta_krr.core.models.allocations import NONE_LITERAL, format_diff, format_recommendation_value
+from robusta_krr.core.models.config import settings
 from robusta_krr.core.models.result import ResourceScan, ResourceType, Result
 
 logger = logging.getLogger("krr")
@@ -52,10 +53,10 @@ def csv_exporter(result: Result) -> str:
     # We need to order the resource columns so that they are in the format of Namespace,Name,Pods,Old Pods,Type,Container,CPU Diff,CPU Requests,CPU Limits,Memory Diff,Memory Requests,Memory Limits
     csv_columns = ["Namespace", "Name", "Pods", "Old Pods", "Type", "Container"]
 
-    if result.config and result.config.show_cluster_name:
+    if settings.show_cluster_name:
         csv_columns.insert(0, "Cluster")
 
-    if result.config and result.config.show_severity:
+    if settings.show_severity:
         csv_columns.append("Severity")
 
     for resource in ResourceType:
diff --git a/robusta_krr/main.py b/robusta_krr/main.py
index a1929f6..f82224b 100644
--- a/robusta_krr/main.py
+++ b/robusta_krr/main.py
@@ -7,6 +7,7 @@
 from typing import List, Optional
 from uuid import UUID
 
+import click
 import typer
 import urllib3
 from pydantic import ValidationError  # noqa: F401
@@ -268,6 +269,8 @@ def run_strategy(
                 **strategy_args,
             ) -> None:
                 f"""Run KRR using the `{_strategy_name}` strategy"""
+                if not show_severity and format != "csv":
+                    raise click.BadOptionUsage("--exclude-severity", "--exclude-severity works only with format=csv")
 
                 try:
                     config = Config(
diff --git a/tests/formatters/test_csv_formatter.py b/tests/formatters/test_csv_formatter.py
index 43841d0..150b2aa 100644
--- a/tests/formatters/test_csv_formatter.py
+++ b/tests/formatters/test_csv_formatter.py
@@ -5,6 +5,7 @@
 
 import pytest
 
+from robusta_krr.core.models.config import Config
 from robusta_krr.core.models.result import Result
 from robusta_krr.formatters.csv import csv_exporter
 
@@ -151,6 +152,7 @@ def _load_result(override_config: dict[str, Any]) -> Result:
     res_data = json.loads(RESULT)
     res_data["config"].update(override_config)
     result = Result(**res_data)
+    Config.set_config(result.config)
     return result
 
 
diff --git a/tests/test_krr.py b/tests/test_krr.py
index d4c6d6a..90ea2af 100644
--- a/tests/test_krr.py
+++ b/tests/test_krr.py
@@ -30,7 +30,7 @@ def test_run(log_flag: str):
         raise e from result.exception
 
 
-@pytest.mark.parametrize("format", ["json", "yaml", "table", "pprint"])
+@pytest.mark.parametrize("format", ["json", "yaml", "table", "pprint", "csv"])
 @pytest.mark.parametrize("output", ["--logtostderr", "-q"])
 def test_output_formats(format: str, output: str):
     result = runner.invoke(app, [STRATEGY_NAME, output, "-f", format])
diff --git a/tests/test_runner.py b/tests/test_runner.py
new file mode 100644
index 0000000..dcf4e19
--- /dev/null
+++ b/tests/test_runner.py
@@ -0,0 +1,21 @@
+import pytest
+from click.testing import Result
+from typer.testing import CliRunner
+
+from robusta_krr.main import app, load_commands
+
+runner = CliRunner(mix_stderr=False)
+load_commands()
+
+
+@pytest.mark.parametrize(
+    "args, expected_exit_code",
+    [
+        (["--exclude-severity", "-f", "csv"], 0),
+        (["--exclude-severity", "-f", "table"], 2),
+        (["--exclude-severity"], 2),
+    ],
+)
+def test_exclude_severity_option(args: list[str], expected_exit_code: int) -> None:
+    result: Result = runner.invoke(app, ["simple", *args])
+    assert result.exit_code == expected_exit_code