import mock import re import socket import threading import time import warnings from unittest import TestCase import pytest from ddtrace.api import API, Response from ddtrace.compat import iteritems, httplib, PY3 from ddtrace.internal.runtime.container import CGroupInfo from ddtrace.vendor.six.moves import BaseHTTPServer, socketserver class _BaseHTTPRequestHandler(BaseHTTPServer.BaseHTTPRequestHandler): error_message_format = '%(message)s\n' error_content_type = 'text/plain' @staticmethod def log_message(format, *args): # noqa: A002 pass class _APIEndpointRequestHandlerTest(_BaseHTTPRequestHandler): def do_PUT(self): self.send_error(200, 'OK') class _TimeoutAPIEndpointRequestHandlerTest(_BaseHTTPRequestHandler): def do_PUT(self): # This server sleeps longer than our timeout time.sleep(5) class _ResetAPIEndpointRequestHandlerTest(_BaseHTTPRequestHandler): def do_PUT(self): return _HOST = '0.0.0.0' _TIMEOUT_PORT = 8743 _RESET_PORT = _TIMEOUT_PORT + 1 class UDSHTTPServer(socketserver.UnixStreamServer, BaseHTTPServer.HTTPServer): def server_bind(self): BaseHTTPServer.HTTPServer.server_bind(self) def _make_uds_server(path, request_handler): server = UDSHTTPServer(path, request_handler) t = threading.Thread(target=server.serve_forever) # Set daemon just in case something fails t.daemon = True t.start() return server, t @pytest.fixture def endpoint_uds_server(tmp_path): server, thread = _make_uds_server(str(tmp_path / 'uds_server_socket'), _APIEndpointRequestHandlerTest) try: yield server finally: server.shutdown() thread.join() def _make_server(port, request_handler): server = BaseHTTPServer.HTTPServer((_HOST, port), request_handler) t = threading.Thread(target=server.serve_forever) # Set daemon just in case something fails t.daemon = True t.start() return server, t @pytest.fixture(scope='module') def endpoint_test_timeout_server(): server, thread = _make_server(_TIMEOUT_PORT, _TimeoutAPIEndpointRequestHandlerTest) try: yield thread finally: server.shutdown() thread.join() @pytest.fixture(scope='module') def endpoint_test_reset_server(): server, thread = _make_server(_RESET_PORT, _ResetAPIEndpointRequestHandlerTest) try: yield thread finally: server.shutdown() thread.join() class ResponseMock: def __init__(self, content, status=200): self.status = status self.content = content def read(self): return self.content def test_api_str(): api = API('localhost', 8126, https=True) assert str(api) == 'https://localhost:8126' api = API('localhost', 8126, '/path/to/uds') assert str(api) == 'unix:///path/to/uds' class APITests(TestCase): def setUp(self): # DEV: Mock here instead of in tests, before we have patched `httplib.HTTPConnection` self.conn = mock.MagicMock(spec=httplib.HTTPConnection) self.api = API('localhost', 8126) def tearDown(self): del self.api del self.conn def test_typecast_port(self): api = API('localhost', u'8126') self.assertEqual(api.port, 8126) @mock.patch('logging.Logger.debug') def test_parse_response_json(self, log): test_cases = { 'OK': dict( js=None, log='Cannot parse Datadog Agent response, please make sure your Datadog Agent is up to date', ), 'OK\n': dict( js=None, log='Cannot parse Datadog Agent response, please make sure your Datadog Agent is up to date', ), 'error:unsupported-endpoint': dict( js=None, log='Unable to parse Datadog Agent JSON response: \'error:unsupported-endpoint\'', ), 42: dict( # int as key to trigger TypeError js=None, log='Unable to parse Datadog Agent JSON response: 42', ), '{}': dict(js={}), '[]': dict(js=[]), # Priority sampling "rate_by_service" response ('{"rate_by_service": ' '{"service:,env:":0.5, "service:mcnulty,env:test":0.9, "service:postgres,env:test":0.6}}'): dict( js=dict( rate_by_service={ 'service:,env:': 0.5, 'service:mcnulty,env:test': 0.9, 'service:postgres,env:test': 0.6, }, ), ), ' [4,2,1] ': dict(js=[4, 2, 1]), } for k, v in iteritems(test_cases): log.reset_mock() r = Response.from_http_response(ResponseMock(k)) js = r.get_json() assert v['js'] == js if 'log' in v: log.assert_called_once() msg = log.call_args[0][0] % log.call_args[0][1:] assert re.match(v['log'], msg), msg @mock.patch('ddtrace.compat.httplib.HTTPConnection') def test_put_connection_close(self, HTTPConnection): """ When calling API._put we close the HTTPConnection we create """ HTTPConnection.return_value = self.conn with warnings.catch_warnings(record=True) as w: self.api._put('/test', '', 1) self.assertEqual(len(w), 0, 'Test raised unexpected warnings: {0!r}'.format(w)) self.conn.request.assert_called_once() self.conn.close.assert_called_once() @mock.patch('ddtrace.compat.httplib.HTTPConnection') def test_put_connection_close_exception(self, HTTPConnection): """ When calling API._put raises an exception we close the HTTPConnection we create """ HTTPConnection.return_value = self.conn # Ensure calling `request` raises an exception self.conn.request.side_effect = Exception with warnings.catch_warnings(record=True) as w: with self.assertRaises(Exception): self.api._put('/test', '', 1) self.assertEqual(len(w), 0, 'Test raised unexpected warnings: {0!r}'.format(w)) self.conn.request.assert_called_once() self.conn.close.assert_called_once() def test_https(): conn = mock.MagicMock(spec=httplib.HTTPSConnection) api = API('localhost', 8126, https=True) with mock.patch('ddtrace.compat.httplib.HTTPSConnection') as HTTPSConnection: HTTPSConnection.return_value = conn api._put('/test', '', 1) conn.request.assert_called_once() conn.close.assert_called_once() def test_flush_connection_timeout_connect(): payload = mock.Mock() payload.get_payload.return_value = 'foobar' payload.length = 12 api = API(_HOST, 2019) response = api._flush(payload) if PY3: assert isinstance(response, (OSError, ConnectionRefusedError)) # noqa: F821 else: assert isinstance(response, socket.error) assert response.errno in (99, 111) def test_flush_connection_timeout(endpoint_test_timeout_server): payload = mock.Mock() payload.get_payload.return_value = 'foobar' payload.length = 12 api = API(_HOST, _TIMEOUT_PORT) response = api._flush(payload) assert isinstance(response, socket.timeout) def test_flush_connection_reset(endpoint_test_reset_server): payload = mock.Mock() payload.get_payload.return_value = 'foobar' payload.length = 12 api = API(_HOST, _RESET_PORT) response = api._flush(payload) if PY3: assert isinstance(response, (httplib.BadStatusLine, ConnectionResetError)) # noqa: F821 else: assert isinstance(response, httplib.BadStatusLine) def test_flush_connection_uds(endpoint_uds_server): payload = mock.Mock() payload.get_payload.return_value = 'foobar' payload.length = 12 api = API(_HOST, 2019, uds_path=endpoint_uds_server.server_address) response = api._flush(payload) assert response.status == 200 @mock.patch('ddtrace.internal.runtime.container.get_container_info') def test_api_container_info(get_container_info): # When we have container information # DEV: `get_container_info` will return a `CGroupInfo` with a `container_id` or `None` info = CGroupInfo(container_id='test-container-id') get_container_info.return_value = info api = API(_HOST, 8126) assert api._container_info is info assert api._headers['Datadog-Container-Id'] == 'test-container-id' # When we do not have container information get_container_info.return_value = None api = API(_HOST, 8126) assert api._container_info is None assert 'Datadog-Container-Id' not in api._headers