test(cli): verify layer selection reaches local and remote execution

This commit is contained in:
Joseph Magly
2026-09-07 16:20:38 -04:00
parent 5fe46d2dca
commit bbf06fd6c8
2 changed files with 34 additions and 5 deletions
+12 -5
View File
@@ -550,7 +550,8 @@ def test_remote_cli_accepts_valid_ssh_port_and_dispatches(monkeypatch):
assert dispatch.call_args.args[0].ssh_port == 2222
def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp_path):
@pytest.mark.parametrize("layer_selection", [None, "all", "middle60"])
def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp_path, layer_selection):
import obliteratus.abliterate
import obliteratus.community
import obliteratus.hard_negative
@@ -604,9 +605,11 @@ def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp
quantization=None, gpu_memory_utilization=0.95, large_model=False,
verify_sample_size=3, refusal_max_tokens=512,
residue_file=["audit.json"], dataset="builtin", residue_weight=2, residue_max=4,
contribute=True, contribute_notes="fixture",
contribute=True, contribute_notes="fixture", layer_selection=layer_selection,
)
cli._cmd_abliterate(args)
assert factory.call_args.kwargs["layer_selection"] == layer_selection
pipeline.run.assert_called_once_with()
assert factory.call_args.kwargs["refusal_max_tokens"] == 512
assert factory.call_args.kwargs["gpu_memory_utilization"] == 0.95
assert factory.call_args.kwargs["trust_remote_code"] is False
@@ -614,7 +617,8 @@ def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp
telemetry.assert_called_once_with(pipeline)
def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path):
@pytest.mark.parametrize("layer_selection", [None, "all", "middle60"])
def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path, layer_selection):
import obliteratus.abliterate
import obliteratus.hard_negative
import obliteratus.model_profile
@@ -651,7 +655,7 @@ def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path):
direction_method="svd", min_layer_fraction=None, max_layer_fraction=None,
harmless_pc_count=None, shield_concept_count=None, shield_ridge=None,
shield_residualize=None, shield_layer_penalty=None, projection_target=None,
device="cpu", dry_run=True,
device="cpu", dry_run=True, layer_selection=layer_selection,
)
cli._cmd_self_improve(args)
assert (output / "self_improve_plan.json").is_file()
@@ -660,7 +664,10 @@ def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path):
result.mkdir()
pipeline = MagicMock()
pipeline.run.return_value = str(result)
monkeypatch.setattr(obliteratus.abliterate, "AbliterationPipeline", Mock(return_value=pipeline))
factory = Mock(return_value=pipeline)
monkeypatch.setattr(obliteratus.abliterate, "AbliterationPipeline", factory)
args.dry_run = False
cli._cmd_self_improve(args)
assert factory.call_args.kwargs["layer_selection"] == layer_selection
pipeline.run.assert_called_once_with()
assert (result / "hard_negative_residue.json").is_file()
+22
View File
@@ -123,3 +123,25 @@ def test_every_offered_strategy_is_one_distill_actually_dispatches_on():
assert branches | {"knee_cosmic"} == EXPECTED_CHOICES, (
f"_distill dispatches on {branches | {'knee_cosmic'}}, CLI offers {EXPECTED_CHOICES}"
)
@pytest.mark.parametrize("strategy", [None, "all", "middle60"])
def test_remote_cli_preserves_override_and_method_default(strategy, monkeypatch):
from unittest.mock import Mock
runner = Mock()
runner.run_obliterate.return_value = "result"
monkeypatch.setattr(cli, "_make_remote_runner", lambda args: runner)
argv = ["obliterate", "org/model", "--remote", "operator@example.invalid",
"--method", "optimized"]
if strategy is not None:
argv += ["--layer-selection", strategy]
cli.main(argv)
forwarded = runner.run_obliterate.call_args.kwargs
assert forwarded["method"] == "optimized"
if strategy is None:
assert "layer_selection" not in forwarded
else:
assert forwarded["layer_selection"] == strategy