diff --git a/datadog_lambda/extension.py b/datadog_lambda/extension.py index c2087a5f..d66848ff 100644 --- a/datadog_lambda/extension.py +++ b/datadog_lambda/extension.py @@ -1,7 +1,15 @@ import logging -import requests 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" @@ -14,7 +22,7 @@ def is_extension_running(): if not path.exists(EXTENSION_PATH): return False try: - requests.get(AGENT_URL + HELLO_PATH) + urllib.request.urlopen(AGENT_URL + HELLO_PATH) except Exception as e: logger.debug("Extension is not running, returned with error %s", e) return False @@ -23,7 +31,8 @@ def is_extension_running(): def flush_extension(): try: - requests.post(AGENT_URL + FLUSH_PATH, data={}) + req = urllib.request.Request(AGENT_URL + FLUSH_PATH, "".encode("ascii")) + urllib.request.urlopen(req) except Exception as e: logger.debug("Failed to flush extension, returned with error %s", e) return False diff --git a/setup.py b/setup.py index cc2fc78b..c8332351 100644 --- a/setup.py +++ b/setup.py @@ -39,6 +39,12 @@ "setuptools>=44.1.1; python_version < '3.0'", ], extras_require={ - "dev": ["nose2==0.9.1", "flake8==3.7.9", "requests==2.22.0", "boto3==1.10.33"] + "dev": [ + "nose2==0.9.1", + "flake8==3.7.9", + "requests==2.22.0", + "boto3==1.10.33", + "httpretty==0.9.7", + ] }, ) diff --git a/tests/test_extension.py b/tests/test_extension.py new file mode 100644 index 00000000..77bf946e --- /dev/null +++ b/tests/test_extension.py @@ -0,0 +1,78 @@ +import os +import sys +import unittest +import httpretty + +try: + from unittest.mock import patch +except ImportError: + from mock import patch + +from datadog_lambda.extension import ( + is_extension_running, + flush_extension, + should_use_extension, +) + + +def exceptionCallback(request, uri, headers): + raise Exception("oopsy!") + + +class TestLambdaExtension(unittest.TestCase): + # do not execute tests for Python v2.x + __test__ = sys.version_info >= (3, 0) + + @patch("datadog_lambda.extension.EXTENSION_PATH", os.path.abspath(__file__)) + def test_is_extension_running_true(self): + httpretty.enable() + last_request = httpretty.last_request() + httpretty.register_uri(httpretty.GET, "http://127.0.0.1:8124/lambda/hello") + assert is_extension_running() == True + assert httpretty.last_request() != last_request + httpretty.disable() + + def test_is_extension_running_file_not_found(self): + httpretty.enable() + last_request = httpretty.last_request() + httpretty.register_uri(httpretty.GET, "http://127.0.0.1:8124/lambda/hello") + assert is_extension_running() == False + assert httpretty.last_request() == last_request + httpretty.disable() + + @patch("datadog_lambda.extension.EXTENSION_PATH", os.path.abspath(__file__)) + def test_is_extension_running_http_failure(self): + httpretty.enable() + last_request = httpretty.last_request() + httpretty.register_uri( + httpretty.GET, + "http://127.0.0.1:8124/lambda/hello", + status=503, + body=exceptionCallback, + ) + assert is_extension_running() == False + assert httpretty.last_request() != last_request + httpretty.disable() + + @patch("datadog_lambda.extension.EXTENSION_PATH", os.path.abspath(__file__)) + def test_flush_ok(self): + httpretty.enable() + last_request = httpretty.last_request() + httpretty.register_uri(httpretty.POST, "http://127.0.0.1:8124/lambda/flush") + assert flush_extension() == True + assert httpretty.last_request() != last_request + httpretty.disable() + + @patch("datadog_lambda.extension.EXTENSION_PATH", os.path.abspath(__file__)) + def test_flush_not_ok(self): + httpretty.enable() + last_request = httpretty.last_request() + httpretty.register_uri( + httpretty.POST, + "http://127.0.0.1:8124/lambda/flush", + status=503, + body=exceptionCallback, + ) + assert flush_extension() == False + assert httpretty.last_request() != last_request + httpretty.disable()