diff --git a/basic_pitch/predict.py b/basic_pitch/predict.py index 2b924e2..ec0978c 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 0000000..8c6e8bf --- /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