diff --git a/datadog_lambda/extension.py b/datadog_lambda/extension.py index d66848ff..159048d7 100644 --- a/datadog_lambda/extension.py +++ b/datadog_lambda/extension.py @@ -1,36 +1,23 @@ 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 +26,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()