From 2db922824e67c812efe227930f9616c2396e287d Mon Sep 17 00:00:00 2001
From: AJ Stuyvenberg <astuyve@gmail.com>
Date: Mon, 4 Mar 2024 13:46:29 -0500
Subject: [PATCH 1/2] feat: Remove http check for extension hello route

---
 datadog_lambda/extension.py | 27 ++++++---------------------
 tests/test_extension.py     | 13 +++----------
 2 files changed, 9 insertions(+), 31 deletions(-)

diff --git a/datadog_lambda/extension.py b/datadog_lambda/extension.py
index d66848ff..e9387b7b 100644
--- a/datadog_lambda/extension.py
+++ b/datadog_lambda/extension.py
@@ -1,36 +1,21 @@
 import logging
 from os import path
 
-try:
-    # only available in python 3
-    # not an issue since the extension is not compatible with python 2.x runtime
-    # https://docs.aws.amazon.com/lambda/latest/dg/using-extensions.html
-    import urllib.request
-except ImportError:
-    # safe since both calls to urllib are protected with try/expect and will return false
-    urllib = None
-
 AGENT_URL = "http://127.0.0.1:8124"
-HELLO_PATH = "/lambda/hello"
 FLUSH_PATH = "/lambda/flush"
 EXTENSION_PATH = "/opt/extensions/datadog-agent"
 
 logger = logging.getLogger(__name__)
 
 
-def is_extension_running():
-    if not path.exists(EXTENSION_PATH):
-        return False
-    try:
-        urllib.request.urlopen(AGENT_URL + HELLO_PATH)
-    except Exception as e:
-        logger.debug("Extension is not running, returned with error %s", e)
-        return False
-    return True
-
+def is_extension_present():
+    if path.exists(EXTENSION_PATH):
+        return True
+    return False
 
 def flush_extension():
     try:
+        import urllib.request
         req = urllib.request.Request(AGENT_URL + FLUSH_PATH, "".encode("ascii"))
         urllib.request.urlopen(req)
     except Exception as e:
@@ -39,4 +24,4 @@ def flush_extension():
     return True
 
 
-should_use_extension = is_extension_running()
+should_use_extension = is_extension_present()
diff --git a/tests/test_extension.py b/tests/test_extension.py
index 5ecb0e36..92142a9e 100644
--- a/tests/test_extension.py
+++ b/tests/test_extension.py
@@ -6,7 +6,7 @@
 from unittest.mock import patch
 
 from datadog_lambda.extension import (
-    is_extension_running,
+    is_extension_present,
     flush_extension,
     should_use_extension,
 )
@@ -48,19 +48,12 @@ def tearDown(self):
 
     @patch("datadog_lambda.extension.EXTENSION_PATH", os.path.abspath(__file__))
     def test_is_extension_running_true(self):
-        assert is_extension_running()
-        assert self.server.called
+        assert is_extension_present()
 
     def test_is_extension_running_file_not_found(self):
-        assert not is_extension_running()
+        assert not is_extension_present()
         assert not self.server.called
 
-    @patch("datadog_lambda.extension.EXTENSION_PATH", os.path.abspath(__file__))
-    def test_is_extension_running_http_failure(self):
-        self.server.raises = True
-        assert not is_extension_running()
-        assert self.server.called
-
     @patch("datadog_lambda.extension.EXTENSION_PATH", os.path.abspath(__file__))
     def test_flush_ok(self):
         assert flush_extension()

From ba1be34e0f09ad883baa35ed82df7ac97f0524c6 Mon Sep 17 00:00:00 2001
From: AJ Stuyvenberg <astuyve@gmail.com>
Date: Mon, 4 Mar 2024 13:51:04 -0500
Subject: [PATCH 2/2] feat: lint

---
 datadog_lambda/extension.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/datadog_lambda/extension.py b/datadog_lambda/extension.py
index e9387b7b..159048d7 100644
--- a/datadog_lambda/extension.py
+++ b/datadog_lambda/extension.py
@@ -13,9 +13,11 @@ def is_extension_present():
         return True
     return False
 
+
 def flush_extension():
     try:
         import urllib.request
+
         req = urllib.request.Request(AGENT_URL + FLUSH_PATH, "".encode("ascii"))
         urllib.request.urlopen(req)
     except Exception as e: