Prepare for requests removal : add 100% unit test coverage on extension.py (#166)

Author: maxday

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

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",
+        ]
     },
 )

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()