diff --git a/apache/tests/conftest.py b/apache/tests/conftest.py index 63c78ecf14d26..aff6051ef838c 100644 --- a/apache/tests/conftest.py +++ b/apache/tests/conftest.py @@ -4,11 +4,11 @@ import os -import mock import pytest import requests from datadog_checks.apache import Apache +from datadog_checks.base.utils.http_testing import MockHTTPResponse from datadog_checks.dev import docker_run from datadog_checks.dev.conditions import CheckEndpoints, WaitFor @@ -48,19 +48,19 @@ def check_status_page_ready(): @pytest.fixture -def mock_hide_server_version(): - req = mock.MagicMock() - with mock.patch('datadog_checks.base.utils.http.requests.Session', return_value=req): - - def mock_requests_get_headers(*args, **kwargs): - r = requests.get(*args, **kwargs) - old_iter = r.iter_lines - r.iter_lines = mock.MagicMock() - r.iter_lines.return_value = (l for l in old_iter(decode_unicode=True) if 'ServerVersion' not in l) - return r - - req.get = mock_requests_get_headers - yield +def mock_hide_server_version(mock_http): + def filter_server_version(url, *args, **kwargs): + r = requests.get(url, **kwargs) + content = '\n'.join(line for line in r.text.splitlines() if 'ServerVersion' not in line) + return MockHTTPResponse( + content=content, + status_code=r.status_code, + headers=dict(r.headers), + url=r.url, + ) + + mock_http.get.side_effect = filter_server_version + yield @pytest.fixture