mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Fix lint: typed image-part filter, restore monkeypatch-based tests
This commit is contained in:
Yucheng Zhu
parent
14816df1c8
commit
49573acf18
2 changed files with 131 additions and 76 deletions
|
|
@ -156,6 +156,31 @@ def _header_value(headers: Mapping[str, str], key: str, default: str) -> str:
|
|||
return headers.get(key, default)
|
||||
|
||||
|
||||
def _is_image_part(item: object) -> bool:
|
||||
"""Whether a structured-message content part carries an image rather than text."""
|
||||
|
||||
if not isinstance(item, Mapping):
|
||||
return False
|
||||
|
||||
part: Final[Mapping[object, object]] = item
|
||||
return part.get("type") == "image_url"
|
||||
|
||||
|
||||
def _scannable_text(content: object) -> str:
|
||||
"""Flatten a structured message's content into the single string the v1 detection endpoint takes.
|
||||
|
||||
Image parts are dropped: the endpoint accepts one string, so an image would only reach it as
|
||||
its stringified source (a base64 blob or a URL), which is not text the scanner can evaluate.
|
||||
"""
|
||||
|
||||
if not isinstance(content, list):
|
||||
return str(content or "")
|
||||
|
||||
parts: Final[Sequence[object]] = content
|
||||
text_parts: Final = [item for item in parts if not _is_image_part(item)] # mutable-ok: sent as a list repr
|
||||
return str(text_parts or "")
|
||||
|
||||
|
||||
def is_saas(host: str) -> bool:
|
||||
"""Checks whether the connection is to the SaaS platform"""
|
||||
|
||||
|
|
@ -263,18 +288,6 @@ class HiddenlayerGuardrail(CustomGuardrail):
|
|||
|
||||
if scan_params := inputs.get("structured_messages"):
|
||||
last_msg: Final = scan_params[-1]
|
||||
content = last_msg.get("content")
|
||||
scan_contents = []
|
||||
|
||||
# Remove images from multimodal content before sending to HiddenLayer v1
|
||||
if content and isinstance(content, list):
|
||||
for item in content:
|
||||
if isinstance(item, dict) and not item.get("type") == "image_url":
|
||||
scan_contents.append(item)
|
||||
effective_content: Final = str(scan_contents) if scan_contents else ""
|
||||
else:
|
||||
effective_content: Final = str(content or "")
|
||||
|
||||
result: _HiddenlayerResponse = await self._call_hiddenlayer(
|
||||
project_id,
|
||||
hl_request_metadata,
|
||||
|
|
@ -282,7 +295,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
|
|||
"messages": [
|
||||
{
|
||||
"role": last_msg.get("role", "user"),
|
||||
"content": effective_content
|
||||
"content": _scannable_text(last_msg.get("content")),
|
||||
}
|
||||
]
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1,34 +1,39 @@
|
|||
import os
|
||||
import sys
|
||||
from typing import cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import List, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from httpx import Response
|
||||
from httpx import Request, Response
|
||||
import requests
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
import litellm
|
||||
from litellm import ModelResponse
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import (
|
||||
HiddenlayerGuardrail,
|
||||
HiddenlayerGuardrailV2,
|
||||
_get_jwt,
|
||||
)
|
||||
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionMessageToolCall,
|
||||
Choices,
|
||||
GenericGuardrailAPIInputs,
|
||||
Message,
|
||||
)
|
||||
|
||||
|
||||
def test_hiddenlayer_config_saas():
|
||||
def test_hiddenlayer_config_saas(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test Hiddenlayer SaaS configuration with init_guardrails_v2."""
|
||||
litellm.set_verbose = True
|
||||
litellm.guardrail_name_config_map = {}
|
||||
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
|
||||
|
||||
# Set environment variables for testing
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
init_guardrails_v2(
|
||||
all_guardrails=[
|
||||
|
|
@ -46,8 +51,6 @@ def test_hiddenlayer_config_saas():
|
|||
)
|
||||
|
||||
# Clean up
|
||||
if "HIDDENLAYER_API_BASE" in os.environ:
|
||||
del os.environ["HIDDENLAYER_API_BASE"]
|
||||
|
||||
|
||||
class TestHiddenlayerGuardrail:
|
||||
|
|
@ -67,9 +70,9 @@ class TestHiddenlayerGuardrail:
|
|||
if key in os.environ:
|
||||
del os.environ[key]
|
||||
|
||||
def test_initialization(self):
|
||||
def test_initialization(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test successful initialization with default values."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -80,19 +83,18 @@ class TestHiddenlayerGuardrail:
|
|||
assert guardrail.guardrail_name == "hiddenlayer"
|
||||
assert guardrail.event_hook == "pre_call"
|
||||
|
||||
def test_initialization_fails_when_api_key_missing(self):
|
||||
def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that initialization fails when API key is not set."""
|
||||
# Ensure API key is not set
|
||||
if "HIDDENLAYER_CLIENT_SECRET" in os.environ:
|
||||
del os.environ["HIDDENLAYER_CLIENT_SECRET"]
|
||||
monkeypatch.delenv("HIDDENLAYER_CLIENT_SECRET", raising=False)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
HiddenlayerGuardrail(guardrail_name="hiddenlayer", event_hook="pre_call")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_no_violations(self):
|
||||
async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for request with no violations detected."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
# Setup guardrail
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
|
|
@ -147,9 +149,9 @@ class TestHiddenlayerGuardrail:
|
|||
assert call_args.args[0] == f"{guardrail.api_base}/detection/v1/interactions"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_with_violations(self):
|
||||
async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for request with violations detected."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
# Setup guardrail
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
|
|
@ -205,9 +207,9 @@ class TestHiddenlayerGuardrail:
|
|||
assert "Blocked by Hiddenlayer" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_no_violations(self):
|
||||
async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for response with no violations detected."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
# Setup guardrail
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
|
|
@ -275,10 +277,10 @@ class TestHiddenlayerGuardrail:
|
|||
mock_post.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_with_violations(self):
|
||||
async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for response with violations detected."""
|
||||
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
# Setup guardrail
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
|
|
@ -344,10 +346,10 @@ class TestHiddenlayerGuardrail:
|
|||
assert exc_info.value.status_code == 400
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_api_error_handling(self):
|
||||
async def test_apply_guardrail_api_error_handling(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test handling of API errors in apply_guardrail."""
|
||||
# Set required API key
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -387,15 +389,17 @@ class TestHiddenlayerGuardrail:
|
|||
assert result == inputs
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_with_call_hiddenlayer_method(self):
|
||||
async def test_validate_with_call_hiddenlayer_method(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test the _validate_with_guard_server internal method."""
|
||||
# Set required API key
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
)
|
||||
|
||||
payload = {"messages": [{"role": "user", "content": "test"}]}
|
||||
|
||||
# Mock successful response
|
||||
mock_response = MagicMock(spec=Response)
|
||||
mock_response.json.return_value = {"evaluation": {"action": "Allow"}}
|
||||
|
|
@ -427,9 +431,9 @@ class TestHiddenlayerGuardrail:
|
|||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_with_image(self):
|
||||
async def test_apply_guardrail_request_with_image(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail strips images from multimodal content before sending to HiddenLayer v1."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -487,17 +491,15 @@ class TestHiddenlayerGuardrail:
|
|||
call_kwargs = mock_post.call_args.kwargs
|
||||
sent_content = call_kwargs["json"]["input"]["messages"][0]["content"]
|
||||
assert isinstance(sent_content, str)
|
||||
assert sent_content == str(
|
||||
[{"type": "text", "text": "how much is on this receipt?"}]
|
||||
)
|
||||
assert sent_content == str([{"type": "text", "text": "how much is on this receipt?"}])
|
||||
|
||||
# Result should be returned without error
|
||||
assert result is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_redact_with_image_content(self):
|
||||
async def test_apply_guardrail_redact_with_image_content(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that REDACT action with multimodal content extracts text properly into inputs['texts']."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrail(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -567,12 +569,11 @@ class TestHiddenlayerGuardrail:
|
|||
assert config_model.__name__ == "HiddenlayerGuardrailConfigModel"
|
||||
|
||||
|
||||
def test_hiddenlayer_config_v2():
|
||||
def test_hiddenlayer_config_v2(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test HiddenLayer V2 configuration with init_guardrails_v2."""
|
||||
litellm.set_verbose = True
|
||||
litellm.guardrail_name_config_map = {}
|
||||
monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
|
||||
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
init_guardrails_v2(
|
||||
all_guardrails=[
|
||||
|
|
@ -590,8 +591,6 @@ def test_hiddenlayer_config_v2():
|
|||
config_file_path="",
|
||||
)
|
||||
|
||||
if "HIDDENLAYER_API_BASE" in os.environ:
|
||||
del os.environ["HIDDENLAYER_API_BASE"]
|
||||
|
||||
|
||||
class TestHiddenlayerGuardrailV2:
|
||||
|
|
@ -609,9 +608,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
if key in os.environ:
|
||||
del os.environ[key]
|
||||
|
||||
def test_initialization(self):
|
||||
def test_initialization(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test successful initialization with default values."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -621,18 +620,17 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert guardrail.guardrail_name == "hiddenlayer"
|
||||
assert guardrail.event_hook == "pre_call"
|
||||
|
||||
def test_initialization_fails_when_api_key_missing(self):
|
||||
def test_initialization_fails_when_api_key_missing(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that initialization fails when API key is not set for SaaS."""
|
||||
if "HIDDENLAYER_CLIENT_SECRET" in os.environ:
|
||||
del os.environ["HIDDENLAYER_CLIENT_SECRET"]
|
||||
monkeypatch.delenv("HIDDENLAYER_CLIENT_SECRET", raising=False)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
HiddenlayerGuardrailV2(guardrail_name="hiddenlayer", event_hook="pre_call")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_no_violations(self):
|
||||
async def test_apply_guardrail_request_no_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for request with no violations detected."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -688,9 +686,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert "detection/v2/request-evaluations" in call_args.args[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_with_violations(self):
|
||||
async def test_apply_guardrail_request_with_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for request with violations detected (block via header)."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -748,9 +746,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert "Blocked by Hiddenlayer" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_no_violations(self):
|
||||
async def test_apply_guardrail_response_no_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for response with no violations detected."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="post_call", default_on=True
|
||||
|
|
@ -813,9 +811,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert "detection/v2/response-evaluations" in call_args.args[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_with_violations(self):
|
||||
async def test_apply_guardrail_response_with_violations(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for response with violations detected (block via header)."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="post_call", default_on=True
|
||||
|
|
@ -860,9 +858,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert "Blocked by Hiddenlayer" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_response_with_tool_calls(self):
|
||||
async def test_apply_guardrail_response_with_tool_calls(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail for response containing tool calls."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="post_call", default_on=True
|
||||
|
|
@ -880,7 +878,7 @@ class TestHiddenlayerGuardrailV2:
|
|||
]
|
||||
|
||||
inputs = GenericGuardrailAPIInputs(
|
||||
tool_calls=cast(list[ChatCompletionMessageToolCall], tool_calls)
|
||||
tool_calls=cast(List[ChatCompletionMessageToolCall], tool_calls)
|
||||
)
|
||||
|
||||
request_data = {
|
||||
|
|
@ -921,9 +919,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert "detection/v2/response-evaluations" in call_args.args[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_hiddenlayer_uses_correct_endpoints(self):
|
||||
async def test_call_hiddenlayer_uses_correct_endpoints(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that _call_hiddenlayer uses the v2 request/response evaluation endpoints."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -956,9 +954,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert "detection/v2/response-evaluations" in mock_post.call_args.args[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_with_image(self):
|
||||
async def test_apply_guardrail_request_with_image(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test apply_guardrail sends multimodal content (image) to HiddenLayer v2."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -1027,9 +1025,9 @@ class TestHiddenlayerGuardrailV2:
|
|||
assert texts == ["how much is on this receipt?"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_request_with_image_multimodal_response(self):
|
||||
async def test_apply_guardrail_request_with_image_multimodal_response(self, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that new_texts extraction handles multimodal content (list) returned by HiddenLayer v2."""
|
||||
os.environ["HIDDENLAYER_API_BASE"] = "https://my.hiddenlayer"
|
||||
monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer")
|
||||
|
||||
guardrail = HiddenlayerGuardrailV2(
|
||||
guardrail_name="hiddenlayer", event_hook="pre_call", default_on=True
|
||||
|
|
@ -1095,3 +1093,47 @@ class TestHiddenlayerGuardrailV2:
|
|||
config_model = HiddenlayerGuardrailV2.get_config_model()
|
||||
assert config_model is not None
|
||||
assert config_model.__name__ == "HiddenlayerGuardrailConfigModel"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hanging_auth_server():
|
||||
"""A server that accepts the connection and never answers, so only a timeout ends the call."""
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from socketserver import ThreadingMixIn
|
||||
|
||||
stop: threading.Event = threading.Event()
|
||||
|
||||
class SilentRequestHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def do_POST(self):
|
||||
stop.wait(timeout=30)
|
||||
|
||||
def log_message(self, format, *args):
|
||||
pass
|
||||
|
||||
class ThreadedServer(ThreadingMixIn, HTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://127.0.0.1:{server.server_port}"
|
||||
finally:
|
||||
stop.set()
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
def test_get_jwt_gives_up_at_the_timeout_instead_of_blocking_the_event_loop(hanging_auth_server):
|
||||
"""
|
||||
`_get_jwt` runs synchronously inside `_call_hiddenlayer`, so an auth host that
|
||||
accepts and never answers used to park the whole worker's event loop.
|
||||
"""
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
_get_jwt(auth_url=hanging_auth_server, api_id="id", api_key="secret", timeout=1)
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue