Further environment variable fixes

awslabs · May 3, 2023 · 7dc45ae · 7dc45ae
1 parent ba4ef0b
commit 7dc45ae
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 6 deletions.
diff --git a/test/test_mqtt5_credentials.py b/test/test_mqtt5_credentials.py
@@ -377,9 +377,18 @@ def sign_function(transform_args, **kwargs):
         callbacks.future_stopped.result(TIMEOUT)
 
         # Set it back to the cached result
-        os.environ["AWS_ACCESS_KEY_ID"] = cache_access_key
-        os.environ["AWS_SECRET_ACCESS_KEY"] = cache_secret_access_key
-        os.environ["AWS_SESSION_TOKEN"] = cache_token
+        if (cache_access_key != None):
+            os.environ["AWS_ACCESS_KEY_ID"] = cache_access_key
+        else:
+            del os.environ["AWS_ACCESS_KEY_ID"]
+        if (cache_secret_access_key != None):
+            os.environ["AWS_SECRET_ACCESS_KEY"] = cache_secret_access_key
+        else:
+            del os.environ["AWS_SECRET_ACCESS_KEY"]
+        if (cache_token != None):
+            os.environ["AWS_SESSION_TOKEN"] = cache_token
+        else:
+            del os.environ["AWS_SESSION_TOKEN"]
 
 
 if __name__ == 'main':

diff --git a/test/test_mqtt_credentials.py b/test/test_mqtt_credentials.py
@@ -296,9 +296,18 @@ def sign_function(transform_args, **kwargs):
         connection.disconnect().result(TIMEOUT)
 
         # Set it back to the cached result
-        os.environ["AWS_ACCESS_KEY_ID"] = cache_access_key
-        os.environ["AWS_SECRET_ACCESS_KEY"] = cache_secret_access_key
-        os.environ["AWS_SESSION_TOKEN"] = cache_token
+        if (cache_access_key != None):
+            os.environ["AWS_ACCESS_KEY_ID"] = cache_access_key
+        else:
+            del os.environ["AWS_ACCESS_KEY_ID"]
+        if (cache_secret_access_key != None):
+            os.environ["AWS_SECRET_ACCESS_KEY"] = cache_secret_access_key
+        else:
+            del os.environ["AWS_SECRET_ACCESS_KEY"]
+        if (cache_token != None):
+            os.environ["AWS_SESSION_TOKEN"] = cache_token
+        else:
+            del os.environ["AWS_SESSION_TOKEN"]
 
 
 if __name__ == 'main':