Add protect mode validation

Author: tomaisthorpe

aikido_zen/__init__.py

@@ -24,6 +24,8 @@
 from aikido_zen.helpers.aikido_disabled_flag_active import aikido_disabled_flag_active
 
 
+VALID_MODES = ("daemon", "daemon_only", "daemon_disabled")
+
 def protect(mode="daemon", token=""):
     """
     Mode can be set to :
@@ -32,6 +34,10 @@ def protect(mode="daemon", token=""):
     - daemon_disabled : This will import sinks/sources but won't start a background process
     Protect user's application
     """
+    if mode not in VALID_MODES:
+        raise ValueError(
+            f"Invalid mode {mode!r}, expected one of {VALID_MODES}. To pass a token, use protect(token=...)"
+        )
     if aikido_disabled_flag_active():
         # Do not run any aikido code when the disabled flag is on
         return

aikido_zen/init_test.py

@@ -17,3 +17,15 @@ def test_protect_with_django(monkeypatch, caplog):
 def test_protect_sets_token():
     aikido_zen.protect(token="MY_TOKEN_1")
     assert get_token_from_env().token == "MY_TOKEN_1"
+
+
+def test_protect_rejects_invalid_mode():
+    with pytest.raises(
+        ValueError, match=r"Invalid mode .*protect\(token=\.\.\.\)"
+    ):
+        aikido_zen.protect("AIK_RUNTIME_some-token-string")
+
+
+@pytest.mark.parametrize("mode", ["daemon", "daemon_only", "daemon_disabled"])
+def test_protect_accepts_valid_modes(mode):
+    aikido_zen.protect(mode=mode)