From 4e24a6a9574cd94f39e1cff887fcafb6f37f32ca Mon Sep 17 00:00:00 2001 From: QuantumGhost Date: Fri, 16 May 2025 17:11:07 +0800 Subject: [PATCH] test(api): Add tests for request & response logging extension --- .../extensions/test_ext_request_logging.py | 265 ++++++++++++++++++ 1 file changed, 265 insertions(+) create mode 100644 api/tests/unit_tests/extensions/test_ext_request_logging.py diff --git a/api/tests/unit_tests/extensions/test_ext_request_logging.py b/api/tests/unit_tests/extensions/test_ext_request_logging.py new file mode 100644 index 0000000000..4e71469bcc --- /dev/null +++ b/api/tests/unit_tests/extensions/test_ext_request_logging.py @@ -0,0 +1,265 @@ +import json +import logging +from unittest import mock + +import pytest +from flask import Flask, Response + +from configs import dify_config +from extensions import ext_request_logging +from extensions.ext_request_logging import _is_content_type_json, _log_request_finished, init_app + + +def test_is_content_type_json(): + """ + Test the _is_content_type_json function. + """ + + assert _is_content_type_json("application/json") is True + # content type header with charset option. + assert _is_content_type_json("application/json; charset=utf-8") is True + # content type header with charset option, in uppercase. + assert _is_content_type_json("APPLICATION/JSON; CHARSET=UTF-8") is True + assert _is_content_type_json("text/html") is False + assert _is_content_type_json("") is False + + +_KEY_NEEDLE = "needle" +_VALUE_NEEDLE = _KEY_NEEDLE[::-1] +_RESPONSE_NEEDLE = "response" + + +def _get_test_app(): + app = Flask(__name__) + + @app.route("/", methods=["GET", "POST"]) + def handler(): + return _RESPONSE_NEEDLE + + return app + + +# NOTE(QuantumGhost): Due to the design of Flask, we need to use monkey patch to write tests. + + +@pytest.fixture +def mock_request_receiver(monkeypatch) -> mock.Mock: + mock_log_request_started = mock.Mock() + monkeypatch.setattr(ext_request_logging, "_log_request_started", mock_log_request_started) + return mock_log_request_started + + +@pytest.fixture +def mock_response_receiver(monkeypatch) -> mock.Mock: + mock_log_request_finished = mock.Mock() + monkeypatch.setattr(ext_request_logging, "_log_request_finished", mock_log_request_finished) + return mock_log_request_finished + + +@pytest.fixture +def mock_logger(monkeypatch) -> logging.Logger: + _logger = mock.MagicMock(spec=logging.Logger) + monkeypatch.setattr(ext_request_logging, "_logger", _logger) + return _logger + + +@pytest.fixture +def enable_request_logging(monkeypatch): + monkeypatch.setattr(dify_config, "ENABLE_REQUEST_LOGGING", True) + + +class TestRequestLoggingExtension: + def test_receiver_should_not_be_invoked_if_configuration_is_disabled( + self, + monkeypatch, + mock_request_receiver, + mock_response_receiver, + ): + monkeypatch.setattr(dify_config, "ENABLE_REQUEST_LOGGING", False) + + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + client.get("/") + + mock_request_receiver.assert_not_called() + mock_response_receiver.assert_not_called() + + def test_receiver_should_be_called_if_enabled( + self, + enable_request_logging, + mock_request_receiver, + mock_response_receiver, + ): + """ + Test the request logging extension with JSON data. + """ + + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + client.post("/", json={_KEY_NEEDLE: _VALUE_NEEDLE}) + + mock_request_receiver.assert_called_once() + mock_response_receiver.assert_called_once() + + +class TestLoggingLevel: + @pytest.mark.usefixtures("enable_request_logging") + def test_logging_should_be_skipped_if_level_is_above_debug(self, enable_request_logging, mock_logger): + mock_logger.isEnabledFor.return_value = False + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + client.post("/", json={_KEY_NEEDLE: _VALUE_NEEDLE}) + mock_logger.debug.assert_not_called() + + +class TestRequestReceiverLogging: + @pytest.mark.usefixtures("enable_request_logging") + def test_non_json_request(self, enable_request_logging, mock_logger, mock_response_receiver): + mock_logger.isEnabledFor.return_value = True + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + client.post("/", data="plain text") + assert mock_logger.debug.call_count == 1 + call_args = mock_logger.debug.call_args[0] + assert "Received Request" in call_args[0] + assert call_args[1] == "POST" + assert call_args[2] == "/" + assert "Request Body" not in call_args[0] + + @pytest.mark.usefixtures("enable_request_logging") + def test_json_request(self, enable_request_logging, mock_logger, mock_response_receiver): + mock_logger.isEnabledFor.return_value = True + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + client.post("/", json={_KEY_NEEDLE: _VALUE_NEEDLE}) + assert mock_logger.debug.call_count == 1 + call_args = mock_logger.debug.call_args[0] + assert "Received Request" in call_args[0] + assert "Request Body" in call_args[0] + assert call_args[1] == "POST" + assert call_args[2] == "/" + assert _KEY_NEEDLE in call_args[3] + + @pytest.mark.usefixtures("enable_request_logging") + def test_json_request_with_empty_body(self, enable_request_logging, mock_logger, mock_response_receiver): + mock_logger.isEnabledFor.return_value = True + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + client.post("/", headers={"Content-Type": "application/json"}) + + assert mock_logger.debug.call_count == 1 + call_args = mock_logger.debug.call_args[0] + assert "Received Request" in call_args[0] + assert "Request Body" not in call_args[0] + assert call_args[1] == "POST" + assert call_args[2] == "/" + + @pytest.mark.usefixtures("enable_request_logging") + def test_json_request_with_invalid_json_as_body(self, enable_request_logging, mock_logger, mock_response_receiver): + mock_logger.isEnabledFor.return_value = True + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + client.post( + "/", + headers={"Content-Type": "application/json"}, + data="{", + ) + assert mock_logger.debug.call_count == 0 + assert mock_logger.exception.call_count == 1 + + exception_call_args = mock_logger.exception.call_args[0] + assert exception_call_args[0] == "Failed to parse JSON request" + + +class TestResponseReceiverLogging: + @pytest.mark.usefixtures("enable_request_logging") + def test_non_json_response(self, enable_request_logging, mock_logger): + mock_logger.isEnabledFor.return_value = True + app = _get_test_app() + response = Response( + "OK", + headers={"Content-Type": "text/plain"}, + ) + _log_request_finished(app, response) + assert mock_logger.debug.call_count == 1 + call_args = mock_logger.debug.call_args[0] + assert "Response" in call_args[0] + assert "200" in call_args[1] + assert call_args[2] == "text/plain" + assert "Response Body" not in call_args[0] + + @pytest.mark.usefixtures("enable_request_logging") + def test_json_response(self, enable_request_logging, mock_logger, mock_response_receiver): + mock_logger.isEnabledFor.return_value = True + app = _get_test_app() + response = Response( + json.dumps({_KEY_NEEDLE: _VALUE_NEEDLE}), + headers={"Content-Type": "application/json"}, + ) + _log_request_finished(app, response) + assert mock_logger.debug.call_count == 1 + call_args = mock_logger.debug.call_args[0] + assert "Response" in call_args[0] + assert "Response Body" in call_args[0] + assert "200" in call_args[1] + assert call_args[2] == "application/json" + assert _KEY_NEEDLE in call_args[3] + + @pytest.mark.usefixtures("enable_request_logging") + def test_json_request_with_invalid_json_as_body(self, enable_request_logging, mock_logger, mock_response_receiver): + mock_logger.isEnabledFor.return_value = True + app = _get_test_app() + + response = Response( + "{", + headers={"Content-Type": "application/json"}, + ) + _log_request_finished(app, response) + assert mock_logger.debug.call_count == 0 + assert mock_logger.exception.call_count == 1 + + exception_call_args = mock_logger.exception.call_args[0] + assert exception_call_args[0] == "Failed to parse JSON response" + + +class TestResponseUnmodified: + def test_when_request_logging_disabled(self): + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + response = client.post( + "/", + headers={"Content-Type": "application/json"}, + data="{", + ) + assert response.text == _RESPONSE_NEEDLE + assert response.status_code == 200 + + @pytest.mark.usefixtures("enable_request_logging") + def test_when_request_logging_enabled(self, enable_request_logging): + app = _get_test_app() + init_app(app) + + with app.test_client() as client: + response = client.post( + "/", + headers={"Content-Type": "application/json"}, + data="{", + ) + assert response.text == _RESPONSE_NEEDLE + assert response.status_code == 200