diff --git a/tests/test_restore_multimodal.py b/tests/test_restore_multimodal.py index aaa39b0..af28477 100644 --- a/tests/test_restore_multimodal.py +++ b/tests/test_restore_multimodal.py @@ -192,3 +192,23 @@ class TestMain: sys.argv = old_argv assert (out_dir / "model.safetensors.index.json").exists() + + def test_cli_dispatch(self, tmp_path): + """Test that 'obliteratus restore-multimodal' dispatches correctly.""" + from obliteratus.cli import main as cli_main + + abl_dir = tmp_path / "abl" + abl_dir.mkdir() + stock_dir = tmp_path / "stock" + stock_dir.mkdir() + out_dir = tmp_path / "out" + + _make_fake_model(abl_dir, "model.", n_layers=1) + _make_fake_model(stock_dir, "model.language_model.", n_layers=1) + + cli_main(["restore-multimodal", + "--abliterated", str(abl_dir), + "--stock", str(stock_dir), + "--output", str(out_dir)]) + + assert (out_dir / "model.safetensors.index.json").exists()