mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-25 06:39:57 +02:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| e752ebaeeb | |||
| 2549194bd1 | |||
| 4c580ea1b8 |
@@ -0,0 +1,2 @@
|
|||||||
|
from:python-pytest-poetry
|
||||||
|
# This file was generated automatically by CodeBeaver based on your repository. Learn how to customize it here: https://docs.codebeaver.ai/configuration/
|
||||||
@@ -0,0 +1,161 @@
|
|||||||
|
import pytest
|
||||||
|
import asyncio
|
||||||
|
from fastapi import FastAPI
|
||||||
|
from asyncio import Queue, Event
|
||||||
|
from agentic_security.core.app import create_app, get_tools_inbox, get_stop_event, get_current_run, set_current_run
|
||||||
|
|
||||||
|
class TestApp:
|
||||||
|
"""Test suite for agentic_security.core.app module."""
|
||||||
|
|
||||||
|
def test_create_app(self):
|
||||||
|
"""Test that create_app returns a FastAPI instance."""
|
||||||
|
app = create_app()
|
||||||
|
assert isinstance(app, FastAPI)
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_tools_inbox(self):
|
||||||
|
"""Test that get_tools_inbox returns the global Queue instance."""
|
||||||
|
queue1 = get_tools_inbox()
|
||||||
|
await queue1.put("test item")
|
||||||
|
queue2 = get_tools_inbox()
|
||||||
|
result = queue2.get_nowait()
|
||||||
|
assert result == "test item"
|
||||||
|
|
||||||
|
def test_get_stop_event(self):
|
||||||
|
"""Test that get_stop_event returns the global Event instance and is not set initially."""
|
||||||
|
event = get_stop_event()
|
||||||
|
assert isinstance(event, Event)
|
||||||
|
assert not event.is_set()
|
||||||
|
|
||||||
|
def test_current_run_initial(self):
|
||||||
|
"""Test that get_current_run returns the global current_run with default values initially."""
|
||||||
|
run = get_current_run()
|
||||||
|
# Default values should be empty strings
|
||||||
|
assert run["spec"] == ""
|
||||||
|
assert run["id"] == ""
|
||||||
|
|
||||||
|
def test_set_current_run(self):
|
||||||
|
"""Test that set_current_run correctly updates current_run."""
|
||||||
|
spec = "test run"
|
||||||
|
result = set_current_run(spec)
|
||||||
|
expected_id = hash(id(spec))
|
||||||
|
# Verify that spec is set correctly
|
||||||
|
assert result["spec"] == spec
|
||||||
|
assert result["id"] == expected_id
|
||||||
|
|
||||||
|
def test_current_run_after_set(self):
|
||||||
|
"""Test that get_current_run returns the updated current_run after set_current_run is called."""
|
||||||
|
spec = "another test run"
|
||||||
|
set_current_run(spec)
|
||||||
|
current = get_current_run()
|
||||||
|
assert current["spec"] == spec
|
||||||
|
assert current["id"] == hash(id(spec))
|
||||||
|
def test_tools_inbox_same_instance(self):
|
||||||
|
"""Test that get_tools_inbox returns the same Queue instance by default."""
|
||||||
|
queue1 = get_tools_inbox()
|
||||||
|
queue2 = get_tools_inbox()
|
||||||
|
assert queue1 is queue2
|
||||||
|
|
||||||
|
def test_stop_event_set(self):
|
||||||
|
"""Test that setting the stop event is reflected in subsequent calls."""
|
||||||
|
event = get_stop_event()
|
||||||
|
event.set() # set the global event
|
||||||
|
# Now, subsequent calls should return the same event which is set.
|
||||||
|
event2 = get_stop_event()
|
||||||
|
assert event2.is_set()
|
||||||
|
|
||||||
|
def test_set_current_run_with_none(self):
|
||||||
|
"""Test that set_current_run handles None as a valid input and updates current_run accordingly."""
|
||||||
|
result = set_current_run(None)
|
||||||
|
expected_id = hash(id(None))
|
||||||
|
assert result["spec"] is None
|
||||||
|
assert result["id"] == expected_id
|
||||||
|
|
||||||
|
def test_multiple_current_run_assignments(self):
|
||||||
|
"""Test multiple assignments to current_run to ensure it always updates correctly."""
|
||||||
|
first_spec = "first run"
|
||||||
|
result1 = set_current_run(first_spec)
|
||||||
|
expected_id1 = hash(id(first_spec))
|
||||||
|
assert result1["spec"] == first_spec
|
||||||
|
assert result1["id"] == expected_id1
|
||||||
|
|
||||||
|
second_spec = "second run"
|
||||||
|
result2 = set_current_run(second_spec)
|
||||||
|
expected_id2 = hash(id(second_spec))
|
||||||
|
assert result2["spec"] == second_spec
|
||||||
|
assert result2["id"] == expected_id2
|
||||||
|
|
||||||
|
current = get_current_run()
|
||||||
|
# The current_run should reflect the latest assignment.
|
||||||
|
assert current["spec"] == second_spec
|
||||||
|
assert current["id"] == expected_id2
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_empty_tools_inbox_exception(self):
|
||||||
|
"""Test that calling get_nowait on an empty tools_inbox raises QueueEmpty."""
|
||||||
|
from asyncio import QueueEmpty
|
||||||
|
queue = get_tools_inbox()
|
||||||
|
# Clear any existing items in the queue
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
queue.get_nowait()
|
||||||
|
except QueueEmpty:
|
||||||
|
break
|
||||||
|
with pytest.raises(QueueEmpty):
|
||||||
|
queue.get_nowait()
|
||||||
|
|
||||||
|
def test_set_current_run_with_dict(self):
|
||||||
|
"""Test that set_current_run correctly handles a dictionary input as spec."""
|
||||||
|
spec = {"key": "value"}
|
||||||
|
result = set_current_run(spec)
|
||||||
|
expected_id = hash(id(spec))
|
||||||
|
assert result["spec"] == spec
|
||||||
|
assert result["id"] == expected_id
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_stop_event_wait(self):
|
||||||
|
"""Test that waiting on the stop event returns once the event is set."""
|
||||||
|
event = get_stop_event()
|
||||||
|
event.clear() # ensure event is not set
|
||||||
|
async def waiter():
|
||||||
|
await event.wait()
|
||||||
|
return True
|
||||||
|
waiter_task = asyncio.create_task(waiter())
|
||||||
|
# Wait a moment to ensure the waiter is pending
|
||||||
|
await asyncio.sleep(0.1)
|
||||||
|
assert not waiter_task.done()
|
||||||
|
event.set()
|
||||||
|
result = await waiter_task
|
||||||
|
assert result is True
|
||||||
|
|
||||||
|
def test_set_current_run_with_int(self):
|
||||||
|
"""Test that set_current_run handles an integer input as spec."""
|
||||||
|
spec = 12345
|
||||||
|
result = set_current_run(spec)
|
||||||
|
expected_id = hash(id(spec))
|
||||||
|
assert result["spec"] == spec
|
||||||
|
assert result["id"] == expected_id
|
||||||
|
|
||||||
|
def test_create_app_routes(self):
|
||||||
|
"""Test that create_app returns a FastAPI instance with default routes available."""
|
||||||
|
app = create_app()
|
||||||
|
paths = [route.path for route in app.routes]
|
||||||
|
# Check that the default OpenAPI route exists
|
||||||
|
assert "/openapi.json" in paths
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_tools_inbox_async_put_get_order(self):
|
||||||
|
"""Test that tools_inbox preserves order when items are added and retrieved asynchronously."""
|
||||||
|
queue = get_tools_inbox()
|
||||||
|
# Clear any existing items in the queue
|
||||||
|
from asyncio import QueueEmpty
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
queue.get_nowait()
|
||||||
|
except QueueEmpty:
|
||||||
|
break
|
||||||
|
items = ["first", "second", "third"]
|
||||||
|
for item in items:
|
||||||
|
await queue.put(item)
|
||||||
|
result_items = []
|
||||||
|
for _ in items:
|
||||||
|
result_items.append(await queue.get())
|
||||||
|
assert result_items == items
|
||||||
@@ -0,0 +1,341 @@
|
|||||||
|
import pytest
|
||||||
|
import base64
|
||||||
|
import httpx
|
||||||
|
import asyncio
|
||||||
|
from agentic_security.http_spec import (
|
||||||
|
LLMSpec,
|
||||||
|
parse_http_spec,
|
||||||
|
escape_special_chars_for_json,
|
||||||
|
encode_image_base64_by_url,
|
||||||
|
encode_audio_base64_by_url,
|
||||||
|
InvalidHTTPSpecError,
|
||||||
|
Modality
|
||||||
|
)
|
||||||
|
|
||||||
|
################################################################################
|
||||||
|
# Tests for agentic_security/http_spec.py
|
||||||
|
################################################################################
|
||||||
|
|
||||||
|
def test_escape_special_chars_for_json():
|
||||||
|
"""Test escaping special characters in a prompt for JSON safety."""
|
||||||
|
prompt = 'Line1\nLine2\t"Quote"\\Backslash'
|
||||||
|
escaped = escape_special_chars_for_json(prompt)
|
||||||
|
assert '\\n' in escaped
|
||||||
|
assert '\\t' in escaped
|
||||||
|
assert '\\"' in escaped
|
||||||
|
assert '\\\\' in escaped
|
||||||
|
|
||||||
|
def test_parse_http_spec_text():
|
||||||
|
"""Test parsing a text HTTP spec without image/audio/files requirements."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\nThis is a prompt: <<PROMPT>>"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
assert llm_spec.method == "POST"
|
||||||
|
assert llm_spec.url == "http://example.com/api"
|
||||||
|
assert llm_spec.headers["Content-Type"] == "application/json"
|
||||||
|
assert "<<PROMPT>>" in llm_spec.body
|
||||||
|
assert not llm_spec.has_files
|
||||||
|
assert not llm_spec.has_image
|
||||||
|
assert not llm_spec.has_audio
|
||||||
|
|
||||||
|
def test_parse_http_spec_files():
|
||||||
|
"""Test parsing a HTTP spec with multipart/form-data header indicating files."""
|
||||||
|
spec = "PUT http://example.com/upload\nContent-Type: multipart/form-data\n\nFile upload test"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
assert llm_spec.has_files
|
||||||
|
|
||||||
|
def test_parse_http_spec_image_audio():
|
||||||
|
"""Test parsing a HTTP spec that requires image and audio via placeholders."""
|
||||||
|
spec = "GET http://example.com/api\nContent-Type: application/json\n\nImage: <<BASE64_IMAGE>> and Audio: <<BASE64_AUDIO>>"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
assert llm_spec.has_image
|
||||||
|
assert llm_spec.has_audio
|
||||||
|
|
||||||
|
def test_encode_image_base64_by_url(monkeypatch):
|
||||||
|
"""Test that image encoding returns the correct base64 string with prefix."""
|
||||||
|
dummy_content = b'test_image'
|
||||||
|
class DummyResponse:
|
||||||
|
def __init__(self, content):
|
||||||
|
self.content = content
|
||||||
|
|
||||||
|
def dummy_get(url):
|
||||||
|
return DummyResponse(dummy_content)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "get", dummy_get)
|
||||||
|
result = encode_image_base64_by_url("http://dummyurl.com/image.jpg")
|
||||||
|
expected = "data:image/jpeg;base64," + base64.b64encode(dummy_content).decode("utf-8")
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
def test_encode_audio_base64_by_url(monkeypatch):
|
||||||
|
"""Test that audio encoding returns the correct base64 string with prefix."""
|
||||||
|
dummy_content = b'test_audio'
|
||||||
|
class DummyResponse:
|
||||||
|
def __init__(self, content):
|
||||||
|
self.content = content
|
||||||
|
|
||||||
|
def dummy_get(url):
|
||||||
|
return DummyResponse(dummy_content)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "get", dummy_get)
|
||||||
|
result = encode_audio_base64_by_url("http://dummyurl.com/audio.mp3")
|
||||||
|
expected = "data:audio/mpeg;base64," + base64.b64encode(dummy_content).decode("utf-8")
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_probe_text(monkeypatch):
|
||||||
|
"""Test the probe function for text modality by replacing <<PROMPT>>."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\n{\"prompt\": \"<<PROMPT>>\"}"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, content, timeout):
|
||||||
|
return httpx.Response(200, text="ok")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
response = await llm_spec.probe("Hello")
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "ok" in response.text
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_probe_with_files(monkeypatch):
|
||||||
|
"""Test that probe correctly branches to _probe_with_files when files are provided."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: multipart/form-data\n\nFile data"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
files = {"file": ("dummy.txt", b"data")}
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, files, timeout):
|
||||||
|
return httpx.Response(200, text="file upload ok")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
response = await llm_spec.probe("Unused", files=files)
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "file upload ok" in response.text
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_verify_image(monkeypatch):
|
||||||
|
"""Test verify method branch for image modality by monkeypatching image encoder."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\n{\"image\": \"<<BASE64_IMAGE>>\"}"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
|
||||||
|
# Replace the image encoder to return a dummy string
|
||||||
|
monkeypatch.setattr("agentic_security.http_spec.encode_image_base64_by_url", lambda url="": "dummy_image")
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, content, timeout):
|
||||||
|
# Check that the dummy image is injected in the content
|
||||||
|
assert "dummy_image" in content
|
||||||
|
return httpx.Response(200, text="image ok")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
response = await llm_spec.verify()
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "image ok" in response.text
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_verify_audio(monkeypatch):
|
||||||
|
"""Test verify method branch for audio modality by monkeypatching audio encoder."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\n{\"audio\": \"<<BASE64_AUDIO>>\"}"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
|
||||||
|
monkeypatch.setattr("agentic_security.http_spec.encode_audio_base64_by_url", lambda url: "dummy_audio")
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, content, timeout):
|
||||||
|
# Ensure that the dummy audio string is present in the request content
|
||||||
|
assert "dummy_audio" in content
|
||||||
|
return httpx.Response(200, text="audio ok")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
response = await llm_spec.verify()
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "audio ok" in response.text
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_verify_files(monkeypatch):
|
||||||
|
"""Test verify method branch for files modality where _probe_with_files is invoked."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: multipart/form-data\n\nFile data"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, files, timeout):
|
||||||
|
return httpx.Response(200, text="files ok")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
response = await llm_spec.verify()
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert "files ok" in response.text
|
||||||
|
|
||||||
|
def test_llm_spec_modality_property():
|
||||||
|
"""Test that the modality property reflects the correct modality."""
|
||||||
|
spec_text = "POST http://example.com/api\nContent-Type: application/json\n\nPrompt: <<PROMPT>>"
|
||||||
|
llm_spec_text = parse_http_spec(spec_text)
|
||||||
|
assert llm_spec_text.modality == Modality.TEXT
|
||||||
|
|
||||||
|
spec_image = "POST http://example.com/api\nContent-Type: application/json\n\nImage: <<BASE64_IMAGE>>"
|
||||||
|
llm_spec_image = parse_http_spec(spec_image)
|
||||||
|
assert llm_spec_image.modality == Modality.IMAGE
|
||||||
|
|
||||||
|
spec_audio = "POST http://example.com/api\nContent-Type: application/json\n\nAudio: <<BASE64_AUDIO>>"
|
||||||
|
llm_spec_audio = parse_http_spec(spec_audio)
|
||||||
|
assert llm_spec_audio.modality == Modality.AUDIO
|
||||||
|
|
||||||
|
def test_from_string_invalid():
|
||||||
|
"""Test that LLMSpec.from_string raises an error for an invalid spec."""
|
||||||
|
invalid_spec = "INVALID_SPEC"
|
||||||
|
with pytest.raises(InvalidHTTPSpecError):
|
||||||
|
LLMSpec.from_string(invalid_spec)
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_validate_missing_files():
|
||||||
|
"""Test that LLMSpec.validate raises a ValueError when files are required but missing."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: multipart/form-data\n\nFile upload test"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
with pytest.raises(ValueError, match="Files are required"):
|
||||||
|
llm_spec.validate("test prompt", "", "", {})
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_validate_missing_image():
|
||||||
|
"""Test that LLMSpec.validate raises a ValueError when an image is required but missing."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\nImage: <<BASE64_IMAGE>>"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
with pytest.raises(ValueError, match="An image is required"):
|
||||||
|
llm_spec.validate("test prompt", "", "dummy_audio", {})
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_validate_missing_audio():
|
||||||
|
"""Test that LLMSpec.validate raises a ValueError when audio is required but missing."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\nAudio: <<BASE64_AUDIO>>"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
with pytest.raises(ValueError, match="Audio is required"):
|
||||||
|
llm_spec.validate("test prompt", "dummy_image", "", {})
|
||||||
|
|
||||||
|
def test_fn_alias(monkeypatch):
|
||||||
|
"""Test that LLMSpec.fn is a functional alias for LLMSpec.probe."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\n{\"prompt\": \"<<PROMPT>>\"}"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
|
||||||
|
# Instead of overriding the instance method, verify the alias at the class level.
|
||||||
|
assert LLMSpec.fn is LLMSpec.probe
|
||||||
|
|
||||||
|
def test_escape_special_chars_no_special():
|
||||||
|
"""Test that the escape function returns the original string if no special characters are present."""
|
||||||
|
prompt = "Simple text without specials"
|
||||||
|
escaped = escape_special_chars_for_json(prompt)
|
||||||
|
assert escaped == "Simple text without specials"
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_probe_text_with_special_chars(monkeypatch):
|
||||||
|
"""Test probe for text modality with special characters in prompt ensuring escaped content."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\n{\"prompt\": \"<<PROMPT>>\"}"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, content, timeout):
|
||||||
|
captured['content'] = content
|
||||||
|
return httpx.Response(200, text="ok")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
test_prompt = 'Hello\nWorld\t"Test"'
|
||||||
|
response = await llm_spec.probe(test_prompt)
|
||||||
|
expected_escaped = escape_special_chars_for_json(test_prompt)
|
||||||
|
assert expected_escaped in captured['content']
|
||||||
|
assert response.status_code == 200
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_verify_both_image_audio(monkeypatch):
|
||||||
|
"""Test verify method when both image and audio placeholders are present.
|
||||||
|
Expect a ValueError because only the image branch is triggered by pattern matching and the missing audio causes validation to fail."""
|
||||||
|
spec = ("POST http://example.com/api\nContent-Type: application/json\n\n"
|
||||||
|
"{\"audio\": \"<<BASE64_AUDIO>>\", \"image\":\"<<BASE64_IMAGE>>\"}")
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
# Monkey patch the image encoder to return a dummy value
|
||||||
|
monkeypatch.setattr("agentic_security.http_spec.encode_image_base64_by_url", lambda url="": "dummy_image")
|
||||||
|
with pytest.raises(ValueError, match="Audio is required"):
|
||||||
|
await llm_spec.verify()
|
||||||
|
|
||||||
|
def test_parse_http_spec_invalid_header_format():
|
||||||
|
"""Test that parse_http_spec raises an error when a header line doesn't have the expected 'key: value' format."""
|
||||||
|
invalid_spec = "GET http://example.com/api\nInvalidHeaderWithoutColon\n\nBody with <<PROMPT>>"
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
parse_http_spec(invalid_spec)
|
||||||
|
|
||||||
|
def test_from_string_valid():
|
||||||
|
"""Test that LLMSpec.from_string returns a valid LLMSpec object when given a proper spec string."""
|
||||||
|
spec = "GET http://example.com/api\nContent-Type: application/json\n\n{ \"prompt\": \"<<PROMPT>>\" }"
|
||||||
|
llm_spec = LLMSpec.from_string(spec)
|
||||||
|
assert llm_spec.method == "GET"
|
||||||
|
assert llm_spec.url == "http://example.com/api"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_parse_http_spec_multiline_body():
|
||||||
|
"""Test parsing an HTTP spec with a multiline body to ensure body concatenation works."""
|
||||||
|
spec = (
|
||||||
|
"PATCH http://example.com/api\n"
|
||||||
|
"Content-Type: application/json\n"
|
||||||
|
"\n"
|
||||||
|
"Line one of body\n"
|
||||||
|
"Line two of body\n"
|
||||||
|
"Line three"
|
||||||
|
)
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
# As implemented, the parser concatenates lines without newline delimiters
|
||||||
|
expected_body = "Line one of bodyLine two of bodyLine three"
|
||||||
|
assert llm_spec.body == expected_body
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_encode_image_default_argument(monkeypatch):
|
||||||
|
"""Test that encode_image_base64_by_url works with its default URL argument."""
|
||||||
|
dummy_content = b'default_image'
|
||||||
|
class DummyResponse:
|
||||||
|
def __init__(self, content):
|
||||||
|
self.content = content
|
||||||
|
|
||||||
|
def dummy_get(url):
|
||||||
|
# check that the default URL (which includes 'fluidicon.png') is used
|
||||||
|
assert "fluidicon.png" in url
|
||||||
|
return DummyResponse(dummy_content)
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx, "get", dummy_get)
|
||||||
|
result = encode_image_base64_by_url()
|
||||||
|
expected = "data:image/jpeg;base64," + base64.b64encode(dummy_content).decode("utf-8")
|
||||||
|
assert result == expected
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_probe_without_prompt_placeholder(monkeypatch):
|
||||||
|
"""Test the probe function when the request body does not include the <<PROMPT>> placeholder."""
|
||||||
|
spec = "POST http://example.com/api\nContent-Type: application/json\n\n{\"message\": \"No placeholder here\"}"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, content, timeout):
|
||||||
|
captured['content'] = content
|
||||||
|
return httpx.Response(200, text="ok without placeholder")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
response = await llm_spec.probe("Ignored prompt")
|
||||||
|
assert "No placeholder here" in captured['content']
|
||||||
|
assert response.status_code == 200
|
||||||
|
|
||||||
|
def test_validate_success():
|
||||||
|
"""Test that LLMSpec.validate does not raise an error when all required data is provided."""
|
||||||
|
# Test case for files: files are provided as required
|
||||||
|
spec_files = "POST http://example.com/api\nContent-Type: multipart/form-data\n\nFile upload"
|
||||||
|
llm_spec_files = parse_http_spec(spec_files)
|
||||||
|
llm_spec_files.validate("some prompt", "dummy_image", "dummy_audio", {"file": ("dummy.txt", b"data")})
|
||||||
|
|
||||||
|
# Test case for image: image is provided as required
|
||||||
|
spec_image = "POST http://example.com/api\nContent-Type: application/json\n\nImage: <<BASE64_IMAGE>>"
|
||||||
|
llm_spec_image = parse_http_spec(spec_image)
|
||||||
|
llm_spec_image.validate("some prompt", "dummy_image", "dummy_audio", {})
|
||||||
|
|
||||||
|
# Test case for audio: audio is provided as required
|
||||||
|
spec_audio = "POST http://example.com/api\nContent-Type: application/json\n\nAudio: <<BASE64_AUDIO>>"
|
||||||
|
llm_spec_audio = parse_http_spec(spec_audio)
|
||||||
|
llm_spec_audio.validate("some prompt", "dummy_image", "dummy_audio", {})
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_probe_invalid_url(monkeypatch):
|
||||||
|
"""Test that probe raises an exception when the HTTP client fails due to an invalid URL."""
|
||||||
|
spec = "GET http://nonexistent_url/api\nContent-Type: application/json\n\n{\"prompt\": \"<<PROMPT>>\"}"
|
||||||
|
llm_spec = parse_http_spec(spec)
|
||||||
|
|
||||||
|
async def dummy_request(self, method, url, headers, content, timeout):
|
||||||
|
raise httpx.RequestError("Invalid URL")
|
||||||
|
|
||||||
|
monkeypatch.setattr(httpx.AsyncClient, "request", dummy_request)
|
||||||
|
with pytest.raises(httpx.RequestError):
|
||||||
|
await llm_spec.probe("Test")
|
||||||
Reference in New Issue
Block a user