Add test cases and change os.environ setting

Author: bitterpanda63

aikido_zen/__init__.py

@@ -1,6 +1,7 @@
 """
 Aggregates from the different modules
 """
+
 import os
 
 # Re-export functions :
@@ -26,12 +27,12 @@ def protect(mode="daemon", token=""):
     - daemon_disabled : This will import sinks/sources but won't start a background process
     Protect user's application
     """
-    if token:
-        os.putenv("AIKIDO_TOKEN", token)
-
     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"