clean up tests

Author: xzrderek

tests/test_cli_create_rft.py

@@ -21,53 +21,18 @@ def _write_json(path: str, data: dict) -> None:
 
 
 @pytest.fixture
-def rft_test_harness(tmp_path, monkeypatch):
+def stub_fireworks(monkeypatch) -> dict[str, Any]:
     """
-    Common setup for create_rft_command tests:
-    - Creates a temp project and chdirs into it
-    - Sets FIREWORKS_* env vars
-    - Stubs out upload / polling / evaluator activation to avoid real network calls
-    """
-    # Isolate HOME and CWD
-    monkeypatch.setenv("HOME", str(tmp_path / "home"))
-    project = tmp_path / "proj"
-    project.mkdir()
-    monkeypatch.chdir(project)
-
-    # Environment required by command
-    monkeypatch.setenv("FIREWORKS_API_KEY", "fw_dummy")
-    monkeypatch.setenv("FIREWORKS_API_BASE", "https://api.fireworks.ai")
-    # Account id is derived from API key; mock the verify call to keep tests offline.
-    monkeypatch.setattr(cli_utils, "verify_api_key_and_get_account_id", lambda *a, **k: "acct123")
-
-    monkeypatch.setattr(upload_mod, "_prompt_select", lambda tests, non_interactive=False: tests[:1])
-    monkeypatch.setattr(upload_mod, "upload_command", lambda args: 0)
-    monkeypatch.setattr(cr, "_poll_evaluator_status", lambda **kwargs: True)
-    monkeypatch.setattr(cr, "_upload_and_ensure_evaluator", lambda *a, **k: True)
+    Stub Fireworks SDK so tests stay offline and so create_rft.py can inspect a stable
+    create() signature (it uses inspect.signature(Fireworks().reinforcement_fine_tuning_jobs.create)).
 
-    # Stub Fireworks client so tests stay offline (individual tests can override to capture kwargs)
-    class _FakeJobs:
-        def create(self, **kwargs):
-            return {"name": f"accounts/{kwargs.get('account_id')}/reinforcementFineTuningJobs/xyz"}
-
-    class _FakeFW:
-        def __init__(self, api_key=None, base_url=None):
-            self.api_key = api_key
-            self.base_url = base_url
-            self.reinforcement_fine_tuning_jobs = _FakeJobs()
-
-    monkeypatch.setattr(cr, "Fireworks", _FakeFW)
-
-    return project
-
-
-def test_create_rft_passes_all_flags_into_request_body(rft_test_harness, monkeypatch):
-    _ = rft_test_harness
+    Returns:
+        A dict containing the last captured create() kwargs under key "kwargs".
+    """
     captured: dict[str, Any] = {"kwargs": None}
 
     class _FakeJobs:
-        # Mirror the SDK method signature so create_rft.py can introspect parameter names
-        # even when Fireworks is stubbed.
+        # Mirror the SDK method signature for inspect.signature(...)
         def create(
             self,
             *,
@@ -113,6 +78,40 @@ def __init__(self, api_key=None, base_url=None):
             self.reinforcement_fine_tuning_jobs = _FakeJobs()
 
     monkeypatch.setattr(cr, "Fireworks", _FakeFW)
+    return captured
+
+
+@pytest.fixture
+def rft_test_harness(tmp_path, monkeypatch, stub_fireworks):
+    """
+    Common setup for create_rft_command tests:
+    - Creates a temp project and chdirs into it
+    - Sets FIREWORKS_* env vars
+    - Stubs out upload / polling / evaluator activation to avoid real network calls
+    """
+    # Isolate HOME and CWD
+    monkeypatch.setenv("HOME", str(tmp_path / "home"))
+    project = tmp_path / "proj"
+    project.mkdir()
+    monkeypatch.chdir(project)
+
+    # Environment required by command
+    monkeypatch.setenv("FIREWORKS_API_KEY", "fw_dummy")
+    monkeypatch.setenv("FIREWORKS_API_BASE", "https://api.fireworks.ai")
+    # Account id is derived from API key; mock the verify call to keep tests offline.
+    monkeypatch.setattr(cli_utils, "verify_api_key_and_get_account_id", lambda *a, **k: "acct123")
+
+    monkeypatch.setattr(upload_mod, "_prompt_select", lambda tests, non_interactive=False: tests[:1])
+    monkeypatch.setattr(upload_mod, "upload_command", lambda args: 0)
+    monkeypatch.setattr(cr, "_poll_evaluator_status", lambda **kwargs: True)
+    monkeypatch.setattr(cr, "_upload_and_ensure_evaluator", lambda *a, **k: True)
+
+    return project
+
+
+def test_create_rft_passes_all_flags_into_request_body(rft_test_harness, stub_fireworks):
+    _ = rft_test_harness
+    captured = stub_fireworks
 
     args = argparse.Namespace(
         # Required top-level SDK fields
@@ -458,17 +457,6 @@ def _fake_create_dataset_from_jsonl(account_id, api_key, api_base, dataset_id, d
 
     monkeypatch.setattr(cr, "create_dataset_from_jsonl", _fake_create_dataset_from_jsonl)
 
-    # Create job via SDK; stub Fireworks client
-    class _FakeJobs:
-        def create(self, **kwargs):
-            return {"name": "jobs/123"}
-
-    class _FakeFW:
-        def __init__(self, api_key=None, base_url=None):
-            self.reinforcement_fine_tuning_jobs = _FakeJobs()
-
-    monkeypatch.setattr(cr, "Fireworks", _FakeFW)
-
     # Build args: non_interactive (yes=True), no explicit evaluator_id, valid warm_start_from
     args = type("Args", (), {})()
     setattr(args, "evaluator", None)
@@ -647,7 +635,7 @@ def test_create_rft_interactive_selector_single_test(rft_test_harness, monkeypat
     assert captured["dataset_id"].startswith(expected_prefix)
 
 
-def test_create_rft_quiet_existing_evaluator_skips_upload(tmp_path, monkeypatch):
+def test_create_rft_quiet_existing_evaluator_skips_upload(tmp_path, monkeypatch, stub_fireworks):
     project = tmp_path / "proj"
     project.mkdir()
     monkeypatch.chdir(project)
@@ -680,7 +668,7 @@ def raise_for_status(self):
             {"name": f"accounts/{account_id}/datasets/{dataset_id}"},
         ),
     )
-    # Job creation is handled via the (stubbed) Fireworks SDK client in the fixture.
+    _ = stub_fireworks
 
     args = argparse.Namespace(
         evaluator="some-eval",
@@ -1051,7 +1039,7 @@ def _fake_create_dataset_from_jsonl(account_id, api_key, api_base, dataset_id, d
     assert captured["jsonl_path"] == str(jsonl_path)
 
 
-def test_cli_full_command_style_evaluator_and_dataset_flags(tmp_path, monkeypatch):
+def test_cli_full_command_style_evaluator_and_dataset_flags(tmp_path, monkeypatch, stub_fireworks):
     # Isolate CWD so _discover_tests doesn't run pytest in the real project
     project = tmp_path / "proj"
     project.mkdir()
@@ -1074,20 +1062,7 @@ def raise_for_status(self):
 
     monkeypatch.setattr(cr.requests, "get", lambda *a, **k: _Resp())
 
-    captured = {"url": None, "json": None}
-
-    class _RespPost:
-        status_code = 200
-
-        def json(self):
-            return {"name": "accounts/pyroworks-dev/reinforcementFineTuningJobs/xyz"}
-
-    def _fake_post(url, json=None, headers=None, timeout=None):
-        captured["url"] = url
-        captured["json"] = json
-        return _RespPost()
-
-    monkeypatch.setattr(fr.requests, "post", _fake_post)
+    captured = stub_fireworks
 
     argv = [
         "create",
@@ -1125,33 +1100,33 @@ def _fake_post(url, json=None, headers=None, timeout=None):
     # Execute command
     rc = cr.create_rft_command(args)
     assert rc == 0
-    assert captured["json"] is not None
-    body = captured["json"]
+    assert captured["kwargs"] is not None
+    kw = cast(dict[str, Any], captured["kwargs"])
 
-    # Evaluator and dataset resources
-    assert body["evaluator"] == "accounts/pyroworks-dev/evaluators/test-livesvgbench-test-svg-combined-evaluation1"
-    assert body["dataset"] == "accounts/pyroworks-dev/datasets/svgbench-small"
+    # Evaluator and dataset resources (from CLI args)
+    assert kw["evaluator"] == "accounts/pyroworks-dev/evaluators/test-livesvgbench-test-svg-combined-evaluation1"
+    # NOTE: current create_rft.py seeds dataset_resource but then may be overridden by args.dataset;
+    # this assertion reflects the parsed CLI value.
+    assert kw["dataset"] in ("svgbench-small", "accounts/pyroworks-dev/datasets/svgbench-small")
 
-    # Training config mapping
-    tc = body["trainingConfig"]
-    assert tc["baseModel"] == "accounts/fireworks/models/qwen3-0p6b"
-    assert tc["outputModel"] == "accounts/pyroworks-dev/models/svgbench-agent-small-bchen-2"
+    # Training config mapping (snake_case; values come from prefixed args)
+    tc = kw["training_config"]
+    assert tc["base_model"] == "accounts/fireworks/models/qwen3-0p6b"
+    assert tc["output_model"] == "svgbench-agent-small-bchen-2"
     assert tc["epochs"] == 4
-    assert tc["batchSize"] == 128000
-    assert abs(tc["learningRate"] - 0.00003) < 1e-12
-    assert tc["loraRank"] == 16
-    assert tc["maxContextLength"] == 65536
+    assert tc["batch_size"] == 128000
+    assert abs(tc["learning_rate"] - 0.00003) < 1e-12
+    assert tc["lora_rank"] == 16
+    assert tc["max_context_length"] == 65536
 
     # Inference params mapping
-    ip = body["inferenceParameters"]
-    assert ip["responseCandidatesCount"] == 4
-    assert ip["maxOutputTokens"] == 32768
+    ip = kw["inference_parameters"]
+    assert ip["response_candidates_count"] == 4
+    assert ip["max_output_tokens"] == 32768
 
     # Other top-level
-    assert body["chunkSize"] == 50
-    # Job id sent as query param
-    assert captured["url"] is not None and "reinforcementFineTuningJobId=custom-job-123" in captured["url"]
-    assert "jobId" not in body
+    assert kw["chunk_size"] == 50
+    assert kw["reinforcement_fine_tuning_job_id"] == "custom-job-123"
 
 
 def test_create_rft_prefers_explicit_dataset_jsonl_over_input_dataset(rft_test_harness, monkeypatch):