From fe78554174ffa9a6b3a7fe59a688c6ff6146b303 Mon Sep 17 00:00:00 2001 From: dissonance Date: Sat, 18 Jul 2026 22:47:13 -0400 Subject: [PATCH] Return nonzero status on CLI failure --- basic_pitch/predict.py | 7 +++++-- tests/test_predict.py | 26 ++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) create mode 100644 tests/test_predict.py diff --git a/basic_pitch/predict.py b/basic_pitch/predict.py index 2b924e26..ec0978c2 100644 --- a/basic_pitch/predict.py +++ b/basic_pitch/predict.py @@ -31,7 +31,7 @@ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" -def main() -> None: +def main() -> int: """Handle command line arguments. Entrypoint for this script.""" parser = argparse.ArgumentParser(description="Predict midi from audio.") parser.add_argument("output_dir", type=str, help="directory to save outputs") @@ -185,14 +185,17 @@ def main() -> None: args.midi_tempo, ) print("\n✨ Done ✨\n") + return 0 except IOError as ioe: print(ioe) + return 1 except Exception as e: print("🚨 Something went wrong 😔 - see the traceback below for details.") print("") print(e) print(traceback.format_exc()) + return 1 if __name__ == "__main__": - main() + raise SystemExit(main()) diff --git a/tests/test_predict.py b/tests/test_predict.py new file mode 100644 index 00000000..8c6e8bf2 --- /dev/null +++ b/tests/test_predict.py @@ -0,0 +1,26 @@ +import sys + +from basic_pitch import inference, predict + + +def run_cli(monkeypatch, tmp_path, predict_and_save): + output_dir = tmp_path / "output" + audio_path = tmp_path / "input.wav" + monkeypatch.setattr(sys, "argv", ["basic-pitch", str(output_dir), str(audio_path)]) + monkeypatch.setattr(inference, "verify_output_dir", lambda path: None) + monkeypatch.setattr(inference, "verify_input_path", lambda path: None) + monkeypatch.setattr(inference, "predict_and_save", predict_and_save) + monkeypatch.setattr(predict, "Model", lambda path: object()) + return predict.main() + + +def test_main_returns_zero_after_success(monkeypatch, tmp_path): + assert run_cli(monkeypatch, tmp_path, lambda *args: None) == 0 + + +def test_main_returns_nonzero_after_inference_failure(monkeypatch, tmp_path, capsys): + def fail(*args): + raise RuntimeError("inference failed") + + assert run_cli(monkeypatch, tmp_path, fail) == 1 + assert "inference failed" in capsys.readouterr().out