Migrate config assertion tests away from requests

airflow/tests/test_unit.py

@@ -26,18 +26,16 @@ def test_service_checks_cannot_connect(aggregator):
     'json_resp, expected_healthy_status, expected_healthy_value',
     [({'status': 'OK'}, AgentCheck.OK, 1), ({'status': 'KO'}, AgentCheck.CRITICAL, 0), ({}, AgentCheck.CRITICAL, 0)],
 )
-def test_service_checks_healthy_exp(aggregator, json_resp, expected_healthy_status, expected_healthy_value):
+def test_service_checks_healthy_exp(aggregator, mock_http, json_resp, expected_healthy_status, expected_healthy_value):
     instance = common.FULL_CONFIG['instances'][0]
     check = AirflowCheck('airflow', common.FULL_CONFIG, [instance])
 
-    with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value=None):
-        mock_session = mock.MagicMock()
-        with mock.patch('datadog_checks.base.utils.http.requests.Session', return_value=mock_session):
-            mock_resp = mock.MagicMock(status_code=200)
-            mock_resp.json.side_effect = [json_resp]
-            mock_session.get.return_value = mock_resp
+    mock_resp = mock.MagicMock(status_code=200)
+    mock_resp.json.side_effect = [json_resp]
+    mock_http.get.return_value = mock_resp
 
-            check.check(None)
+    with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value=None):
+        check.check(None)
 
     tags = ['key:my-tag', 'url:http://localhost:8080']
 
@@ -54,65 +52,60 @@ def test_service_checks_healthy_exp(aggregator, json_resp, expected_healthy_stat
     ],
 )
 def test_service_checks_healthy_stable(
-    aggregator, metadb_status, scheduler_status, expected_healthy_status, expected_healthy_value
+    aggregator, mock_http, metadb_status, scheduler_status, expected_healthy_status, expected_healthy_value
 ):  # Stable is only defined in the context of Airflow 2
     instance = common.FULL_CONFIG['instances'][0]
     check = AirflowCheck('airflow', common.FULL_CONFIG, [instance])
 
-    with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value='2.6.2'):
-        mock_session = mock.MagicMock()
-        with mock.patch('datadog_checks.base.utils.http.requests.Session', return_value=mock_session):
-            mock_resp = mock.MagicMock(status_code=200)
-            mock_resp.json.side_effect = [
-                {'metadatabase': {'status': metadb_status}, 'scheduler': {'status': scheduler_status}},
-                {'status': 'OK'},
-            ]
-            mock_session.get.return_value = mock_resp
+    mock_resp = mock.MagicMock(status_code=200)
+    mock_resp.json.side_effect = [
+        {'metadatabase': {'status': metadb_status}, 'scheduler': {'status': scheduler_status}},
+        {'status': 'OK'},
+    ]
+    mock_http.get.return_value = mock_resp
 
-            check.check(None)
+    with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value='2.6.2'):
+        check.check(None)
 
     tags = ['key:my-tag', 'url:http://localhost:8080']
 
     aggregator.assert_service_check('airflow.healthy', expected_healthy_status, tags=tags, count=1)
     aggregator.assert_metric('airflow.healthy', expected_healthy_value, tags=tags, count=1)
 
 
-def test_dag_total_tasks(aggregator, task_instance):
+def test_dag_total_tasks(aggregator, mock_http, task_instance):
     instance = common.FULL_CONFIG['instances'][0]
     check = AirflowCheck('airflow', common.FULL_CONFIG, [instance])
 
-    with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value='2.6.2'):
-        req = mock.MagicMock()
-        with mock.patch('datadog_checks.base.utils.http.requests.Session', return_value=req):
-            mock_resp = mock.MagicMock(status_code=200)
-            mock_resp.json.side_effect = [
-                {'metadatabase': {'status': 'healthy'}, 'scheduler': {'status': 'healthy'}},
-                task_instance,
-            ]
-            req.get.return_value = mock_resp
+    mock_resp = mock.MagicMock(status_code=200)
+    mock_resp.json.side_effect = [
+        {'metadatabase': {'status': 'healthy'}, 'scheduler': {'status': 'healthy'}},
+        task_instance,
+    ]
+    mock_http.get.return_value = mock_resp
 
-            check.check(None)
+    with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value='2.6.2'):
+        check.check(None)
 
     aggregator.assert_metric('airflow.dag.task.total_running', value=1, count=1)
 
 
-def test_dag_task_ongoing_duration(aggregator, task_instance):
+def test_dag_task_ongoing_duration(aggregator, mock_http, task_instance):
     instance = common.FULL_CONFIG['instances'][0]
     check = AirflowCheck('airflow', common.FULL_CONFIG, [instance])
 
+    mock_resp = mock.MagicMock(status_code=200)
+    mock_resp.json.side_effect = [
+        {'metadatabase': {'status': 'healthy'}, 'scheduler': {'status': 'healthy'}},
+    ]
+    mock_http.get.return_value = mock_resp
+
     with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value='2.6.2'):
-        req = mock.MagicMock()
-        with mock.patch('datadog_checks.base.utils.http.requests.Session', return_value=req):
-            mock_resp = mock.MagicMock(status_code=200)
-            mock_resp.json.side_effect = [
-                {'metadatabase': {'status': 'healthy'}, 'scheduler': {'status': 'healthy'}},
-            ]
-            req.get.return_value = mock_resp
-            with mock.patch(
-                'datadog_checks.airflow.airflow.AirflowCheck._get_all_task_instances',
-                return_value=task_instance.get('task_instances'),
-            ):
-                check.check(None)
+        with mock.patch(
+            'datadog_checks.airflow.airflow.AirflowCheck._get_all_task_instances',
+            return_value=task_instance.get('task_instances'),
+        ):
+            check.check(None)
 
     aggregator.assert_metric(
         'airflow.dag.task.ongoing_duration',
@@ -141,23 +134,21 @@ def test_dag_task_ongoing_duration(aggregator, task_instance):
         ),
     ],
 )
-def test_config_collect_ongoing_duration(collect_ongoing_duration, should_call_method):
+def test_config_collect_ongoing_duration(mock_http, collect_ongoing_duration, should_call_method):
     instance = {**common.FULL_CONFIG['instances'][0], 'collect_ongoing_duration': collect_ongoing_duration}
     check = AirflowCheck('airflow', common.FULL_CONFIG, [instance])
 
+    mock_resp = mock.MagicMock(status_code=200)
+    mock_resp.json.side_effect = [
+        {'metadatabase': {'status': 'healthy'}, 'scheduler': {'status': 'healthy'}},
+    ]
+    mock_http.get.return_value = mock_resp
+
     with mock.patch('datadog_checks.airflow.airflow.AirflowCheck._get_version', return_value='2.6.2'):
-        req = mock.MagicMock()
-        with mock.patch('datadog_checks.base.utils.http.requests.Session', return_value=req):
-            mock_resp = mock.MagicMock(status_code=200)
-            mock_resp.json.side_effect = [
-                {'metadatabase': {'status': 'healthy'}, 'scheduler': {'status': 'healthy'}},
-            ]
-            req.get.return_value = mock_resp
-
-            with mock.patch(
-                'datadog_checks.airflow.airflow.AirflowCheck._get_all_task_instances'
-            ) as mock_get_all_task_instances:
-                check.check(None)
-
-                # Assert method calls
-                mock_get_all_task_instances.assert_has_calls(should_call_method, any_order=False)
+        with mock.patch(
+            'datadog_checks.airflow.airflow.AirflowCheck._get_all_task_instances'
+        ) as mock_get_all_task_instances:
+            check.check(None)
+
+            # Assert method calls
+            mock_get_all_task_instances.assert_has_calls(should_call_method, any_order=False)

consul/tests/test_unit.py

@@ -3,7 +3,6 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import logging
 
-import mock
 import pytest
 
 from datadog_checks.consul import ConsulCheck
@@ -155,25 +154,25 @@ def test_get_nodes_with_service_critical(aggregator):
     aggregator.assert_metric('consul.catalog.services_count', value=1, tags=expected_tags)
 
 
-def test_consul_request(aggregator, instance, mocker):
+def test_consul_request(aggregator, instance, mocker, mock_http):
     consul_check = ConsulCheck(common.CHECK_NAME, {}, [consul_mocks.MOCK_CONFIG])
     mocker.patch("datadog_checks.base.utils.serialization.json.loads")
-    with mock.patch("datadog_checks.consul.consul.requests.Session.get") as mock_requests_get:
+
+    consul_check.consul_request("foo")
+    url = "{}/{}".format(instance["url"], "foo")
+    aggregator.assert_service_check("consul.can_connect", ConsulCheck.OK, tags=["url:{}".format(url)], count=1)
+
+    aggregator.reset()
+    mock_http.get.side_effect = Exception("message")
+    with pytest.raises(Exception):
         consul_check.consul_request("foo")
-        url = "{}/{}".format(instance["url"], "foo")
-        aggregator.assert_service_check("consul.can_connect", ConsulCheck.OK, tags=["url:{}".format(url)], count=1)
-
-        aggregator.reset()
-        mock_requests_get.side_effect = Exception("message")
-        with pytest.raises(Exception):
-            consul_check.consul_request("foo")
-        aggregator.assert_service_check(
-            "consul.can_connect",
-            ConsulCheck.CRITICAL,
-            tags=["url:{}".format(url)],
-            count=1,
-            message="Consul request to {} failed: message".format(url),
-        )
+    aggregator.assert_service_check(
+        "consul.can_connect",
+        ConsulCheck.CRITICAL,
+        tags=["url:{}".format(url)],
+        count=1,
+        message="Consul request to {} failed: message".format(url),
+    )
 
 
 def test_service_checks(aggregator):
@@ -648,26 +647,13 @@ def test_network_latency_node_name(
         ),
     ],
 )
-def test_config(test_case, extra_config, expected_http_kwargs, mocker):
+def test_config(test_case, extra_config, expected_http_kwargs):
     instance = extra_config
     check = ConsulCheck(common.CHECK_NAME, {}, instances=[instance])
-    mocker.patch("datadog_checks.base.utils.serialization.json.loads")
 
-    with mock.patch('datadog_checks.base.utils.http.requests.Session') as session:
-        mock_session = mock.MagicMock()
-        session.return_value = mock_session
-        mock_session.get.return_value = mock.MagicMock(status_code=200)
-
-        check.check(None)
-
-        http_wargs = {
-            'auth': mock.ANY,
-            'cert': mock.ANY,
-            'headers': mock.ANY,
-            'proxies': mock.ANY,
-            'timeout': mock.ANY,
-            'verify': mock.ANY,
-            'allow_redirects': mock.ANY,
-        }
-        http_wargs.update(expected_http_kwargs)
-        mock_session.get.assert_called_with('/v1/status/leader', **http_wargs)
+    for key, value in expected_http_kwargs.items():
+        if key == 'headers':
+            for h_key, h_value in value.items():
+                assert check.http.get_header(h_key) == h_value
+        else:
+            assert check.http.options[key] == value

couch/tests/test_unit.py

@@ -4,7 +4,6 @@
 from copy import deepcopy
 from unittest.mock import MagicMock
 
-import mock
 import pytest
 
 from datadog_checks.couch import CouchDb
@@ -30,24 +29,8 @@ def test_config(test_case, extra_config, expected_http_kwargs):
     instance.update(extra_config)
     check = CouchDb(common.CHECK_NAME, {}, instances=[instance])
 
-    r = mock.MagicMock()
-    with mock.patch('datadog_checks.base.utils.http.requests.Session', return_value=r):
-        r.get.return_value = mock.MagicMock(status_code=200, content='{}')
-
-        check.check(instance)
-
-        http_wargs = {
-            'auth': mock.ANY,
-            'cert': mock.ANY,
-            'headers': mock.ANY,
-            'proxies': mock.ANY,
-            'timeout': mock.ANY,
-            'verify': mock.ANY,
-            'allow_redirects': mock.ANY,
-        }
-        http_wargs.update(expected_http_kwargs)
-
-        r.get.assert_called_with('http://{}:5984/_all_dbs/'.format(common.HOST), **http_wargs)
+    for key, value in expected_http_kwargs.items():
+        assert check.http.options[key] == value
 
 
 def test_new_version_system_metrics(load_test_data):

datadog_checks_base/datadog_checks/base/utils/http.py

@@ -448,6 +448,20 @@ def __init__(self, instance, init_config, remapper=None, logger=None, session=No
         self.tls_config = {key: value for key, value in config.items() if key.startswith('tls_')}
         self._https_adapters = {}
 
+    def get_header(self, name: str, default: str | None = None) -> str | None:
+        """Look up a request header by name. Lookup is case-insensitive."""
+        for key, value in self.options['headers'].items():
+            if key.lower() == name.lower():
+                return value
+        return default
+
+    def set_header(self, name: str, value: str) -> None:
+        for key in self.options['headers']:
+            if key.lower() == name.lower():
+                self.options['headers'][key] = value
+                return
+        self.options['headers'][name] = value
+
     def get(self, url, **options):
         return self._request('get', url, options)
 

datadog_checks_base/datadog_checks/base/utils/http_protocol.py

@@ -3,10 +3,19 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 from __future__ import annotations
 
+from collections.abc import Mapping
 from typing import Any, Iterator, Protocol
 
 
 class HTTPResponseProtocol(Protocol):
+    status_code: int
+    content: bytes
+    text: str
+    headers: Mapping[str, str]
+
+    def json(self, **kwargs: Any) -> Any: ...
+    def raise_for_status(self) -> None: ...
+    def close(self) -> None: ...
     def iter_content(self, chunk_size: int | None = None, decode_unicode: bool = False) -> Iterator[bytes | str]: ...
     def iter_lines(
         self,
@@ -19,10 +28,14 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool | None: ...
 
 
 class HTTPClientProtocol(Protocol):
+    options: dict[str, Any]
+
     def get(self, url: str, **options: Any) -> HTTPResponseProtocol: ...
     def post(self, url: str, **options: Any) -> HTTPResponseProtocol: ...
     def head(self, url: str, **options: Any) -> HTTPResponseProtocol: ...
     def put(self, url: str, **options: Any) -> HTTPResponseProtocol: ...
     def patch(self, url: str, **options: Any) -> HTTPResponseProtocol: ...
     def delete(self, url: str, **options: Any) -> HTTPResponseProtocol: ...
     def options_method(self, url: str, **options: Any) -> HTTPResponseProtocol: ...
+    def get_header(self, name: str, default: str | None = None) -> str | None: ...
+    def set_header(self, name: str, value: str) -> None: ...

datadog_checks_base/tests/base/utils/http/test_headers.py

@@ -97,3 +97,34 @@ def test_extra_headers_on_http_method_call():
     # make sure the original headers are not modified
     assert http.options['headers'] == complete_headers
     assert extra_headers == {"foo": "bar"}
+
+
+def test_get_header_default_for_missing():
+    http = RequestsWrapper({}, {})
+    assert http.get_header('X-Missing') is None
+    assert http.get_header('X-Missing', 'fallback') == 'fallback'
+
+
+def test_get_header_case_insensitive():
+    http = RequestsWrapper({}, {})
+    assert http.get_header('accept') == '*/*'
+    assert http.get_header('Accept') == '*/*'
+    assert http.get_header('ACCEPT') == '*/*'
+
+
+def test_set_header():
+    http = RequestsWrapper({}, {})
+    http.set_header('X-Token', 'abc123')
+    assert http.get_header('X-Token') == 'abc123'
+    http.set_header('Accept', 'application/json')
+    assert http.get_header('Accept') == 'application/json'
+
+
+def test_set_header_case_insensitive():
+    http = RequestsWrapper({}, {})
+    http.set_header('accept', 'application/json')
+    # Overwrites the existing 'Accept' key (preserving original casing)
+    assert http.get_header('Accept') == 'application/json'
+    assert http.options['headers']['Accept'] == 'application/json'
+    # No duplicate key created
+    assert sum(1 for k in http.options['headers'] if k.lower() == 'accept') == 1