"""Behavior tests for resumable auto-obliteration orchestration.""" from __future__ import annotations import json from pathlib import Path import pytest import obliteratus.abliterate from obliteratus.auto_obliterate import ( AutoObliterateResult, AutoObliterator, IterationResult, ) pytestmark = pytest.mark.cpu def _finish(generator): yielded = [] while True: try: yielded.append(next(generator)) except StopIteration as completed: return yielded, completed.value class _SuccessfulPipeline: created = [] def __init__(self, **kwargs): self.kwargs = kwargs self._quality_metrics = { "perplexity": 4.25, "coherence": 0.9, "refusal_rate": 0.04, "kl_divergence": 0.02, } self._strong_layers = [0, 1] self._expert_directions = {0: {0: object(), 1: object()}} self.handle = type("Handle", (), {"model": object(), "tokenizer": object()})() self.created.append(self) def run(self): Path(self.kwargs["output_dir"]).mkdir(parents=True) self.kwargs["on_log"]("pipeline ran") return Path(self.kwargs["output_dir"]) def test_result_round_trip_ignores_unknown_forward_compatible_fields(): original = AutoObliterateResult( model_id="local/model", iterations=[ IterationResult( iteration=1, method="aggressive", prompt_volume=8, categories_targeted=["test"], ), ], success=True, ) encoded = original.to_dict() encoded["future_field"] = "ignored" encoded["iterations"][0]["future_field"] = "ignored" restored = AutoObliterateResult.from_dict(encoded) assert restored == original assert isinstance(restored.iterations[0], IterationResult) def test_valid_state_resumes_at_next_iteration(tmp_path): output = tmp_path / "state" output.mkdir() state = AutoObliterateResult( model_id="local/model", iterations=[IterationResult(1, "aggressive", 8, output_dir="first")], ) (output / "auto_state.json").write_text(json.dumps(state.to_dict())) auto = AutoObliterator("local/model", output_base=str(output), max_iterations=2) assert auto._resume_from == 1 assert auto._result.iterations[0].method == "aggressive" @pytest.mark.parametrize( "invalid_state", [ "not json", "[]", '{"model_id": "different/model"}', '{"model_id": "local/model", "iterations": [5]}', ], ) def test_invalid_state_is_quarantined_with_actionable_warning( tmp_path, caplog, invalid_state, ): output = tmp_path / "state" output.mkdir() state_file = output / "auto_state.json" state_file.write_text(invalid_state) auto = AutoObliterator("local/model", output_base=str(output)) quarantined = list(output.glob("auto_state.json.corrupt-*")) assert auto._resume_from == 0 assert not state_file.exists() assert len(quarantined) == 1 assert quarantined[0].read_text() == invalid_state assert "quarantined at" in caplog.text def test_invalid_state_is_retained_when_quarantine_fails(tmp_path, monkeypatch, caplog): output = tmp_path / "state" output.mkdir() state_file = output / "auto_state.json" state_file.write_text("not json", encoding="utf-8") def fail_quarantine(_source, _target): raise OSError("read-only filesystem") monkeypatch.setattr("obliteratus.auto_obliterate.os.replace", fail_quarantine) auto = AutoObliterator("local/model", output_base=str(output)) assert auto._resume_from == 0 assert state_file.read_text() == "not json" assert "could not be quarantined" in caplog.text assert "read-only filesystem" in caplog.text def test_interrupted_state_replace_preserves_previous_checkpoint( tmp_path, monkeypatch, caplog, ): output = tmp_path / "state" output.mkdir() state_file = output / "auto_state.json" state_file.write_text('{"sentinel": true}', encoding="utf-8") auto = AutoObliterator.__new__(AutoObliterator) auto.model_id = "local/model" auto.output_base = str(output) auto._state_file = state_file auto._result = AutoObliterateResult(model_id="local/model") def fail_replace(_self, _target): raise OSError("simulated interrupted replace") monkeypatch.setattr(Path, "replace", fail_replace) auto._save_state() assert json.loads(state_file.read_text()) == {"sentinel": True} assert not (output / "auto_state.tmp").exists() assert "simulated interrupted replace" in caplog.text def test_auto_loop_runs_pipeline_persists_metrics_and_stops_at_target( tmp_path, monkeypatch, ): _SuccessfulPipeline.created.clear() monkeypatch.setattr( obliteratus.abliterate, "AbliterationPipeline", _SuccessfulPipeline, ) monkeypatch.setattr( AutoObliterator, "_get_expanded_prompts", staticmethod(lambda _iteration: (["harmful"] * 3, ["harmless"] * 3)), ) auto = AutoObliterator( "local/model", output_base=str(tmp_path / "run"), max_iterations=3, target_refusal_rate=0.05, trust_remote_code=False, ) yielded, result = _finish(auto.run()) assert result.success is True assert result.final_refusal_rate == 0.04 assert len(result.iterations) == 1 assert result.iterations[0].strong_layers == 2 assert result.iterations[0].ega_expert_dirs == 2 assert result.final_output_dir.endswith("iter_1") assert len(yielded) == 5 assert yielded[-1][0] == "✅ Complete" created = _SuccessfulPipeline.created[0] assert created.kwargs["model_name"] == "local/model" assert created.kwargs["method"] == "aggressive" assert created.kwargs["harmful_prompts"] == ["harmful"] * 3 saved = json.loads((tmp_path / "run" / "auto_state.json").read_text()) assert saved["success"] is True def test_auto_loop_records_failures_and_completes_without_success( tmp_path, monkeypatch, ): class FailingPipeline: def __init__(self, **_kwargs): pass def run(self): raise RuntimeError("simulated pipeline failure") monkeypatch.setattr(obliteratus.abliterate, "AbliterationPipeline", FailingPipeline) monkeypatch.setattr( AutoObliterator, "_get_expanded_prompts", staticmethod(lambda _iteration: (["harmful"], ["harmless"])), ) auto = AutoObliterator( "local/model", output_base=str(tmp_path / "run"), max_iterations=1, ) yielded, result = _finish(auto.run()) assert result.success is False assert result.final_output_dir == "" assert result.iterations[0].error == "simulated pipeline failure" assert yielded[-1][0] == "⚠️ Complete (target not met)" assert "simulated pipeline" in yielded[-1][2] def test_prompt_expansion_and_benchmark_fallbacks(monkeypatch): monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) assert AutoObliterator._quick_benchmark_claude("missing", "model") == { "method": "skipped", "reason": "no OPENROUTER_API_KEY", } pipeline = type( "Pipeline", (), {"_quality_metrics": {"refusal_rate": 0.2, "coherence": 0.8}}, )() assert AutoObliterator._quick_benchmark_heuristic(pipeline) == { "refusal_rate": 0.2, "perplexity": None, "coherence": 0.8, "kl_divergence": None, "method": "heuristic", } def test_reset_clears_persisted_state(tmp_path): auto = AutoObliterator("local/model", output_base=str(tmp_path / "run")) auto._result.iterations.append(IterationResult(1, "aggressive", 1)) auto._save_state() auto.reset() assert auto._resume_from == 0 assert auto._result.iterations == [] assert not auto._state_file.exists()