Return nonzero status on CLI inference failure

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())

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