From 01e9d5f85b1189847e097164f235387f36198bda Mon Sep 17 00:00:00 2001
From: Mouse <mouseparker123@gmail.com>
Date: Thu, 16 Apr 2026 01:31:50 -0700
Subject: [PATCH] feat: harden CLI config parsing and overrides

---
 promptlens/cli.py                | 48 ++++++++++++++++++++++++++++----
 tests/test_cli_config_loading.py | 47 +++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_cli_config_loading.py

diff --git a/promptlens/cli.py b/promptlens/cli.py
index 2d34ce3..bc051ae 100644
--- a/promptlens/cli.py
+++ b/promptlens/cli.py
@@ -26,6 +26,46 @@
 console = Console()
 
 
+def _load_config_data(config_path: str) -> dict:
+    """Load and validate raw YAML config data."""
+    with open(config_path, "r") as f:
+        config_data = yaml.safe_load(f)
+
+    if config_data is None:
+        raise click.ClickException("Configuration file is empty")
+
+    if not isinstance(config_data, dict):
+        raise click.ClickException(
+            "Configuration root must be a mapping/object"
+        )
+
+    return config_data
+
+
+def _apply_cli_overrides(
+    config_data: dict,
+    golden_set: Optional[str],
+    output_dir: Optional[str],
+) -> dict:
+    """Apply CLI override flags to config data safely."""
+    if golden_set:
+        config_data["golden_set"] = golden_set
+
+    if output_dir:
+        output_config = config_data.get("output")
+        if output_config is None:
+            output_config = {}
+            config_data["output"] = output_config
+        elif not isinstance(output_config, dict):
+            raise click.ClickException(
+                "Invalid configuration: 'output' must be an object"
+            )
+
+        output_config["directory"] = output_dir
+
+    return config_data
+
+
 def setup_logging(level: str = "INFO") -> None:
     """Set up logging configuration.
 
@@ -90,14 +130,10 @@ def run(
     try:
         # Load config
         console.print(f"\n[cyan]Loading configuration from {config}...[/cyan]")
-        with open(config, "r") as f:
-            config_data = yaml.safe_load(f)
+        config_data = _load_config_data(config)
 
         # Override with CLI options
-        if golden_set:
-            config_data["golden_set"] = golden_set
-        if output_dir:
-            config_data["output"]["directory"] = output_dir
+        config_data = _apply_cli_overrides(config_data, golden_set, output_dir)
 
         # Parse config
         try:
diff --git a/tests/test_cli_config_loading.py b/tests/test_cli_config_loading.py
new file mode 100644
index 0000000..382a459
--- /dev/null
+++ b/tests/test_cli_config_loading.py
@@ -0,0 +1,47 @@
+import click
+import pytest
+
+from promptlens.cli import _apply_cli_overrides, _load_config_data
+
+
+def test_load_config_data_rejects_empty_file(tmp_path):
+    config_path = tmp_path / "empty.yaml"
+    config_path.write_text("")
+
+    with pytest.raises(click.ClickException, match="Configuration file is empty"):
+        _load_config_data(str(config_path))
+
+
+def test_load_config_data_rejects_non_mapping_root(tmp_path):
+    config_path = tmp_path / "list_root.yaml"
+    config_path.write_text("- just\n- a\n- list\n")
+
+    with pytest.raises(
+        click.ClickException, match="Configuration root must be a mapping/object"
+    ):
+        _load_config_data(str(config_path))
+
+
+def test_apply_cli_overrides_creates_output_mapping_when_missing():
+    config_data = {"golden_set": "tests.yaml", "models": []}
+
+    updated = _apply_cli_overrides(
+        config_data=config_data,
+        golden_set=None,
+        output_dir="./custom_results",
+    )
+
+    assert updated["output"]["directory"] == "./custom_results"
+
+
+def test_apply_cli_overrides_rejects_non_mapping_output():
+    config_data = {"golden_set": "tests.yaml", "models": [], "output": "bad"}
+
+    with pytest.raises(
+        click.ClickException, match="Invalid configuration: 'output' must be an object"
+    ):
+        _apply_cli_overrides(
+            config_data=config_data,
+            golden_set=None,
+            output_dir="./custom_results",
+        )