Add Test To Verify NEW_RELIC_* Env Vars Not Reset During Upgrade (#159)

* Add Test To Verify NEW_RELIC_* Env Vars Not Reset During Upgrade

Closes #158

* Only remove environment variables the CLI manages

Author: kolanos

newrelic_lambda_cli/layers.py

@@ -11,6 +11,16 @@
 from newrelic_lambda_cli.integrations import _get_license_key_policy_arn
 from newrelic_lambda_cli.types import LayerInstall, LayerUninstall
 
+NEW_RELIC_ENV_VARS = (
+    "NEW_RELIC_ACCOUNT_ID",
+    "NEW_RELIC_EXTENSION_SEND_FUNCTION_LOGS",
+    "NEW_RELIC_LAMBDA_EXTENSION_ENABLED",
+    "NEW_RELIC_LAMBDA_HANDLER",
+    "NEW_RELIC_LICENSE_KEY",
+    "NEW_RELIC_LOG_ENDPOINT",
+    "NEW_RELIC_TELEMETRY_ENDPOINT",
+)
+
 
 def index(region, runtime):
     req = requests.get(
@@ -120,14 +130,9 @@ def _add_new_relic(input, config, nr_license_key):
             "NEW_RELIC_LAMBDA_EXTENSION_ENABLED"
         ] = "true"
 
-        if input.enable_extension_function_logs:
-            update_kwargs["Environment"]["Variables"][
-                "NEW_RELIC_EXTENSION_SEND_FUNCTION_LOGS"
-            ] = "true"
-        else:
-            update_kwargs["Environment"]["Variables"][
-                "NEW_RELIC_EXTENSION_SEND_FUNCTION_LOGS"
-            ] = "false"
+        update_kwargs["Environment"]["Variables"][
+            "NEW_RELIC_EXTENSION_SEND_FUNCTION_LOGS"
+        ] = ("true" if input.enable_extension_function_logs else "false")
 
         if input.nr_region == "staging":
             update_kwargs["Environment"]["Variables"][
@@ -242,7 +247,7 @@ def _remove_new_relic(input, config):
         .get("Environment", {})
         .get("Variables", {})
         .items()
-        if not key.startswith("NEW_RELIC")
+        if key not in NEW_RELIC_ENV_VARS
     }
 
     # Remove New Relic layers

tests/test_layers.py

@@ -152,6 +152,24 @@ def test_add_new_relic(aws_credentials, mock_function_config):
         in config["Configuration"]["Environment"]["Variables"]
     )
 
+    config = mock_function_config("python3.6")
+    config["Configuration"]["Environment"]["Variables"]["NEW_RELIC_FOO"] = "bar"
+    config["Configuration"]["Layers"] = [{"Arn": get_arn_prefix("us-east-1")}]
+    update_kwargs = _add_new_relic(
+        layer_install(
+            session=session,
+            aws_region="us-east-1",
+            nr_account_id=12345,
+            enable_extension=True,
+            enable_extension_function_logs=True,
+            upgrade=True,
+        ),
+        config,
+        "foobarbaz",
+    )
+    assert "NEW_RELIC_FOO" in update_kwargs["Environment"]["Variables"]
+    assert update_kwargs["Layers"][0] != get_arn_prefix("us-east-1")
+
 
 @mock_lambda
 def test_remove_new_relic(aws_credentials, mock_function_config):