Ensure --num-variations is from 1 to 4

Author: msabramo

firefly/cli.py

@@ -123,6 +123,9 @@ def generate(
         raise typer.BadParameter("client_id must be provided as an option or via the FIREFLY_CLIENT_ID environment variable.")
     if not client_secret:
         raise typer.BadParameter("client_secret must be provided as an option or via the FIREFLY_CLIENT_SECRET environment variable.")
+    if num_variations is not None and not (1 <= num_variations <= 4):
+        typer.secho("--num-variations must be between 1 and 4", fg=typer.colors.RED, err=True)
+        raise typer.Exit(code=-1)
     # Parse JSON for style/structure if provided
     style_obj = None
     if style:

tests/test_cli.py

@@ -219,3 +219,33 @@ def test_generate_invalid_json_structure(monkeypatch):
     )
     assert result.exit_code == 2
     assert "Invalid JSON for --structure" in result.output
+
+def test_generate_invalid_num_variations(monkeypatch):
+    # Test too low
+    result = runner.invoke(
+        app,
+        [
+            "image", "generate",
+            "--client-id", "dummy_id",
+            "--client-secret", "dummy_secret",
+            "--prompt", "test",
+            "--num-variations", "0",
+            "--use-mocks"
+        ]
+    )
+    assert result.exit_code == -1
+    assert "--num-variations must be between 1 and 4" in result.output
+    # Test too high
+    result = runner.invoke(
+        app,
+        [
+            "image", "generate",
+            "--client-id", "dummy_id",
+            "--client-secret", "dummy_secret",
+            "--prompt", "test",
+            "--num-variations", "5",
+            "--use-mocks"
+        ]
+    )
+    assert result.exit_code == -1
+    assert "--num-variations must be between 1 and 4" in result.output