Allow setting of token in aikido_zen.protect()

aikido_zen/__init__.py

@@ -2,6 +2,8 @@
 Aggregates from the different modules
 """
 
+import os
+
 # Re-export functions :
 from aikido_zen.context.users import set_user
 from aikido_zen.middleware import should_block_request
@@ -17,7 +19,7 @@
 from aikido_zen.helpers.aikido_disabled_flag_active import aikido_disabled_flag_active
 
 
-def protect(mode="daemon"):
+def protect(mode="daemon", token=""):
     """
     Mode can be set to :
     - daemon : Default, imports sinks/sources and starts background_process
@@ -28,6 +30,9 @@ def protect(mode="daemon"):
     if aikido_disabled_flag_active():
         # Do not run any aikido code when the disabled flag is on
         return
+    if token:
+        os.environ["AIKIDO_TOKEN"] = token
+
     if mode in ("daemon", "daemon_only"):
         start_background_process()
     if mode == "daemon_only":

aikido_zen/init_test.py

@@ -1,6 +1,9 @@
 import pytest
+
+import aikido_zen
 from aikido_zen import protect
 from aikido_zen.background_process import get_comms, reset_comms
+from aikido_zen.helpers.token import get_token_from_env
 
 
 def test_protect_with_django(monkeypatch, caplog):
@@ -9,3 +12,8 @@ def test_protect_with_django(monkeypatch, caplog):
     assert "starting" in caplog.text
     reset_comms()
     assert get_comms() == None
+
+
+def test_protect_sets_token():
+    aikido_zen.protect(token="MY_TOKEN_1")
+    assert get_token_from_env().token == "MY_TOKEN_1"