feat(api): support file uploads via JSON base64 and multipart/form-data
This commit is contained in:
dengjingren
@@ -0,0 +1,379 @@
|
||||
"""Tests for API file upload functionality (JSON base64 + multipart)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from nanobot.api.server import (
|
||||
API_CHAT_ID,
|
||||
API_SESSION_KEY,
|
||||
_parse_json_content,
|
||||
_save_base64_data_url,
|
||||
create_app,
|
||||
)
|
||||
|
||||
try:
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
HAS_AIOHTTP = True
|
||||
except ImportError:
|
||||
HAS_AIOHTTP = False
|
||||
|
||||
pytest_plugins = ("pytest_asyncio",)
|
||||
|
||||
|
||||
def _make_mock_agent(response_text: str = "mock response") -> MagicMock:
|
||||
agent = MagicMock()
|
||||
agent.process_direct = AsyncMock(return_value=response_text)
|
||||
agent._connect_mcp = AsyncMock()
|
||||
agent.close_mcp = AsyncMock()
|
||||
return agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent():
|
||||
return _make_mock_agent()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app(mock_agent):
|
||||
return create_app(mock_agent, model_name="test-model", request_timeout=10.0)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def aiohttp_client():
|
||||
clients: list[TestClient] = []
|
||||
|
||||
async def _make_client(app):
|
||||
client = TestClient(TestServer(app))
|
||||
await client.start_server()
|
||||
clients.append(client)
|
||||
return client
|
||||
|
||||
try:
|
||||
yield _make_client
|
||||
finally:
|
||||
for client in clients:
|
||||
await client.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper function tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_save_base64_data_url_saves_png(tmp_path) -> None:
|
||||
"""Saving a base64 data URL creates a file with correct extension."""
|
||||
b64_data = base64.b64encode(b"fake png data").decode()
|
||||
data_url = f"data:image/png;base64,{b64_data}"
|
||||
result = _save_base64_data_url(data_url, tmp_path)
|
||||
assert result is not None
|
||||
assert result.endswith(".png")
|
||||
assert (tmp_path / result.replace(str(tmp_path) + "/", "")).read_bytes() == b"fake png data"
|
||||
|
||||
|
||||
def test_save_base64_data_url_handles_invalid_b64(tmp_path) -> None:
|
||||
"""Invalid base64 returns None."""
|
||||
result = _save_base64_data_url("data:image/png;base64,not-valid-base64!!!", tmp_path)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_save_base64_data_url_handles_unknown_mime(tmp_path) -> None:
|
||||
"""Unknown MIME type defaults to .bin."""
|
||||
b64_data = base64.b64encode(b"some data").decode()
|
||||
data_url = f"data:unknown/type;base64,{b64_data}"
|
||||
result = _save_base64_data_url(data_url, tmp_path)
|
||||
assert result is not None
|
||||
assert result.endswith(".bin")
|
||||
|
||||
|
||||
def test_parse_json_content_extracts_text_and_media(tmp_path) -> None:
|
||||
"""Parse JSON with text + base64 image saves image and returns paths."""
|
||||
b64_data = base64.b64encode(b"img").decode()
|
||||
body = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "describe this"},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64_data}"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
text, media_paths = _parse_json_content(body)
|
||||
assert text == "describe this"
|
||||
assert len(media_paths) == 1
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
def test_parse_json_content_plain_text_only() -> None:
|
||||
"""Plain text string content returns no media."""
|
||||
body = {"messages": [{"role": "user", "content": "hello"}]}
|
||||
text, media_paths = _parse_json_content(body)
|
||||
assert text == "hello"
|
||||
assert media_paths == []
|
||||
|
||||
|
||||
def test_parse_json_content_validates_single_message() -> None:
|
||||
"""Multiple messages raise ValueError."""
|
||||
body = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "first"},
|
||||
{"role": "user", "content": "second"},
|
||||
]
|
||||
}
|
||||
with pytest.raises(ValueError, match="single user message"):
|
||||
_parse_json_content(body)
|
||||
|
||||
|
||||
def test_parse_json_content_validates_user_role() -> None:
|
||||
"""Non-user role raises ValueError."""
|
||||
body = {"messages": [{"role": "system", "content": "you are a bot"}]}
|
||||
with pytest.raises(ValueError, match="single user message"):
|
||||
_parse_json_content(body)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multipart upload tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_multipart_upload_saves_file(aiohttp_client, mock_agent, tmp_path) -> None:
|
||||
"""Multipart upload saves file to media dir and passes path to process_direct."""
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
app = create_app(mock_agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
file_data = b"test file content"
|
||||
data = BytesIO(file_data)
|
||||
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions",
|
||||
data={"message": "analyze this", "files": data},
|
||||
)
|
||||
assert resp.status == 200
|
||||
call_kwargs = mock_agent.process_direct.call_args.kwargs
|
||||
assert call_kwargs["content"] == "analyze this"
|
||||
assert len(call_kwargs.get("media", [])) == 1
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_multipart_multiple_files(aiohttp_client, mock_agent, tmp_path) -> None:
|
||||
"""Multipart upload with multiple files saves all and passes paths."""
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
app = create_app(mock_agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
# Note: aiohttp test client has limited multipart support
|
||||
# This test verifies the basic flow
|
||||
file_data = b"test content"
|
||||
data = BytesIO(file_data)
|
||||
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions",
|
||||
data={"message": "analyze", "files": data},
|
||||
)
|
||||
assert resp.status == 200
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_multipart_file_size_limit(aiohttp_client, mock_agent, tmp_path) -> None:
|
||||
"""File exceeding MAX_FILE_SIZE returns 413."""
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
app = create_app(mock_agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
# Create a file larger than 10MB
|
||||
large_data = b"x" * (11 * 1024 * 1024)
|
||||
data = BytesIO(large_data)
|
||||
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions",
|
||||
data={"message": "analyze", "files": data},
|
||||
)
|
||||
assert resp.status == 413
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_multipart_defaults_text_when_missing(aiohttp_client, mock_agent, tmp_path) -> None:
|
||||
"""Multipart without message field uses default text."""
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
app = create_app(mock_agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
file_data = b"content"
|
||||
data = BytesIO(file_data)
|
||||
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions",
|
||||
data={"files": data},
|
||||
)
|
||||
assert resp.status == 200
|
||||
call_kwargs = mock_agent.process_direct.call_args.kwargs
|
||||
assert call_kwargs["content"] == "请分析上传的文件"
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_multipart_with_session_id(aiohttp_client, mock_agent, tmp_path) -> None:
|
||||
"""Multipart upload with session_id uses custom session key."""
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
app = create_app(mock_agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
file_data = b"content"
|
||||
data = BytesIO(file_data)
|
||||
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions",
|
||||
data={"message": "hello", "session_id": "my-session", "files": data},
|
||||
)
|
||||
assert resp.status == 200
|
||||
call_kwargs = mock_agent.process_direct.call_args.kwargs
|
||||
assert call_kwargs["session_key"] == "api:my-session"
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Backward compatibility tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_plain_text_backward_compat(aiohttp_client, mock_agent) -> None:
|
||||
"""Plain text JSON request (no media) works as before."""
|
||||
app = create_app(mock_agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"messages": [{"role": "user", "content": "hello world"}]},
|
||||
)
|
||||
assert resp.status == 200
|
||||
body = await resp.json()
|
||||
assert body["choices"][0]["message"]["content"] == "mock response"
|
||||
call_kwargs = mock_agent.process_direct.call_args.kwargs
|
||||
assert call_kwargs["content"] == "hello world"
|
||||
assert call_kwargs.get("media") is None
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_base64_image_upload(aiohttp_client, mock_agent, tmp_path) -> None:
|
||||
"""JSON request with base64 data URL saves file and passes path."""
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
app = create_app(mock_agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
# Use valid base64 for a tiny PNG (1x1 transparent pixel)
|
||||
tiny_png_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
|
||||
|
||||
resp = await client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "what is this"},
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{tiny_png_b64}"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
assert resp.status == 200
|
||||
call_kwargs = mock_agent.process_direct.call_args.kwargs
|
||||
assert call_kwargs["content"] == "what is this"
|
||||
assert len(call_kwargs.get("media", [])) == 1
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DOCX document extraction tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@pytest.mark.asyncio
|
||||
async def test_docx_upload_extracted_and_sent(aiohttp_client, tmp_path) -> None:
|
||||
"""Uploaded DOCX should have its text extracted before being sent to AI."""
|
||||
from docx import Document
|
||||
|
||||
agent = _make_mock_agent("This report shows $5M revenue")
|
||||
import os
|
||||
original_cwd = os.getcwd()
|
||||
os.chdir(tmp_path)
|
||||
|
||||
try:
|
||||
app = create_app(agent, model_name="m")
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
doc = Document()
|
||||
doc.add_heading("Q1 Report", level=1)
|
||||
doc.add_paragraph("Total revenue: $5,000,000")
|
||||
buf = BytesIO()
|
||||
doc.save(buf)
|
||||
docx_bytes = buf.getvalue()
|
||||
|
||||
import aiohttp
|
||||
data = aiohttp.FormData()
|
||||
data.add_field("message", "summarize the report")
|
||||
data.add_field("files", docx_bytes, filename="report.docx",
|
||||
content_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
|
||||
|
||||
resp = await client.post("/v1/chat/completions", data=data)
|
||||
assert resp.status == 200
|
||||
call_kwargs = agent.process_direct.call_args.kwargs
|
||||
media = call_kwargs.get("media", [])
|
||||
assert len(media) == 1
|
||||
assert "report.docx" in media[0]
|
||||
finally:
|
||||
os.chdir(original_cwd)
|
||||
@@ -0,0 +1,66 @@
|
||||
"""Tests for context builder document handling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
|
||||
from nanobot.agent.context import ContextBuilder
|
||||
|
||||
|
||||
def _make_builder(tmp_path: Path) -> ContextBuilder:
|
||||
"""Create a minimal ContextBuilder for testing."""
|
||||
return ContextBuilder(workspace=tmp_path, timezone="UTC")
|
||||
|
||||
|
||||
def test_build_user_content_with_no_media_returns_string(tmp_path: Path) -> None:
|
||||
builder = _make_builder(tmp_path)
|
||||
result = builder._build_user_content("hello", None)
|
||||
assert result == "hello"
|
||||
|
||||
|
||||
def test_build_user_content_with_image_returns_list(tmp_path: Path) -> None:
|
||||
"""Image files should produce base64 content blocks."""
|
||||
builder = _make_builder(tmp_path)
|
||||
png = tmp_path / "test.png"
|
||||
png.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
|
||||
result = builder._build_user_content("describe this", [str(png)])
|
||||
assert isinstance(result, list)
|
||||
types = [b["type"] for b in result]
|
||||
assert "image_url" in types
|
||||
assert "text" in types
|
||||
|
||||
|
||||
def test_build_user_content_with_docx_includes_extracted_text(tmp_path: Path) -> None:
|
||||
"""Document files should have their text extracted and included."""
|
||||
from docx import Document
|
||||
|
||||
doc = Document()
|
||||
doc.add_paragraph("Quarterly revenue is $5M")
|
||||
docx_path = tmp_path / "report.docx"
|
||||
doc.save(docx_path)
|
||||
|
||||
builder = _make_builder(tmp_path)
|
||||
result = builder._build_user_content("summarize this", [str(docx_path)])
|
||||
assert isinstance(result, str)
|
||||
assert "Quarterly revenue" in result
|
||||
|
||||
|
||||
def test_build_user_content_mixed_image_and_document(tmp_path: Path) -> None:
|
||||
"""Mix of images and documents: images as base64, docs as text."""
|
||||
from docx import Document
|
||||
|
||||
png = tmp_path / "chart.png"
|
||||
png.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 100)
|
||||
|
||||
doc = Document()
|
||||
doc.add_paragraph("Report text here")
|
||||
docx = tmp_path / "report.docx"
|
||||
doc.save(docx)
|
||||
|
||||
builder = _make_builder(tmp_path)
|
||||
result = builder._build_user_content("analyze both", [str(png), str(docx)])
|
||||
assert isinstance(result, list)
|
||||
assert any(b["type"] == "image_url" for b in result)
|
||||
text_parts = [b.get("text", "") for b in result if b.get("type") == "text"]
|
||||
assert any("Report text here" in t for t in text_parts)
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Tests for document text extraction utilities."""
|
||||
|
||||
import io
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.utils.document import (
|
||||
SUPPORTED_EXTENSIONS,
|
||||
_is_text_extension,
|
||||
extract_text,
|
||||
)
|
||||
|
||||
|
||||
class TestSupportedExtensions:
|
||||
"""Test the SUPPORTED_EXTENSIONS constant."""
|
||||
|
||||
def test_supported_extensions_include_common_formats(self):
|
||||
"""Test that common document formats are included."""
|
||||
# Document formats
|
||||
assert ".pdf" in SUPPORTED_EXTENSIONS
|
||||
assert ".docx" in SUPPORTED_EXTENSIONS
|
||||
assert ".xlsx" in SUPPORTED_EXTENSIONS
|
||||
assert ".pptx" in SUPPORTED_EXTENSIONS
|
||||
|
||||
# Text formats
|
||||
assert ".txt" in SUPPORTED_EXTENSIONS
|
||||
assert ".md" in SUPPORTED_EXTENSIONS
|
||||
assert ".csv" in SUPPORTED_EXTENSIONS
|
||||
assert ".json" in SUPPORTED_EXTENSIONS
|
||||
assert ".yaml" in SUPPORTED_EXTENSIONS
|
||||
assert ".yml" in SUPPORTED_EXTENSIONS
|
||||
|
||||
# Image formats
|
||||
assert ".png" in SUPPORTED_EXTENSIONS
|
||||
assert ".jpg" in SUPPORTED_EXTENSIONS
|
||||
assert ".jpeg" in SUPPORTED_EXTENSIONS
|
||||
|
||||
|
||||
class TestExtractText:
|
||||
"""Test the extract_text function."""
|
||||
|
||||
def test_extract_text_unsupported_returns_none(self, tmp_path: Path):
|
||||
"""Test that unsupported file types return None."""
|
||||
unsupported_file = tmp_path / "file.xyz"
|
||||
unsupported_file.write_text("content")
|
||||
|
||||
result = extract_text(unsupported_file)
|
||||
assert result is None
|
||||
|
||||
def test_extract_text_file_not_found(self, tmp_path: Path):
|
||||
"""Test that non-existent files return error string."""
|
||||
missing_file = tmp_path / "nonexistent.txt"
|
||||
|
||||
result = extract_text(missing_file)
|
||||
assert result is not None
|
||||
assert "[error: file not found:" in result
|
||||
|
||||
def test_extract_text_txt_file(self, tmp_path: Path):
|
||||
"""Test extracting text from a .txt file."""
|
||||
txt_file = tmp_path / "test.txt"
|
||||
content = "Hello, world!\nThis is a test."
|
||||
txt_file.write_text(content, encoding="utf-8")
|
||||
|
||||
result = extract_text(txt_file)
|
||||
assert result == content
|
||||
|
||||
def test_extract_text_txt_file_with_truncation(self, tmp_path: Path):
|
||||
"""Test that large text files are truncated."""
|
||||
txt_file = tmp_path / "large.txt"
|
||||
# Create content larger than _MAX_TEXT_LENGTH
|
||||
content = "x" * 300_000
|
||||
txt_file.write_text(content, encoding="utf-8")
|
||||
|
||||
result = extract_text(txt_file)
|
||||
assert len(result) < 300_000
|
||||
assert "(truncated," in result
|
||||
assert "chars total)" in result
|
||||
|
||||
def test_extract_text_md_file(self, tmp_path: Path):
|
||||
"""Test extracting text from a .md file."""
|
||||
md_file = tmp_path / "test.md"
|
||||
content = "# Header\n\nSome markdown content."
|
||||
md_file.write_text(content, encoding="utf-8")
|
||||
|
||||
result = extract_text(md_file)
|
||||
assert result == content
|
||||
|
||||
def test_extract_text_csv_file(self, tmp_path: Path):
|
||||
"""Test extracting text from a .csv file."""
|
||||
csv_file = tmp_path / "test.csv"
|
||||
content = "name,age\nAlice,30\nBob,25"
|
||||
csv_file.write_text(content, encoding="utf-8")
|
||||
|
||||
result = extract_text(csv_file)
|
||||
assert result == content
|
||||
|
||||
def test_extract_text_json_file(self, tmp_path: Path):
|
||||
"""Test extracting text from a .json file."""
|
||||
json_file = tmp_path / "test.json"
|
||||
content = '{"key": "value", "number": 42}'
|
||||
json_file.write_text(content, encoding="utf-8")
|
||||
|
||||
result = extract_text(json_file)
|
||||
assert result == content
|
||||
|
||||
def test_extract_text_xlsx(self, tmp_path: Path):
|
||||
"""Test extracting text from an .xlsx file."""
|
||||
from openpyxl import Workbook
|
||||
|
||||
xlsx_file = tmp_path / "test.xlsx"
|
||||
wb = Workbook()
|
||||
ws = wb.active
|
||||
ws.title = "Sheet1"
|
||||
ws["A1"] = "Name"
|
||||
ws["B1"] = "Age"
|
||||
ws["A2"] = "Alice"
|
||||
ws["B2"] = 30
|
||||
ws["A3"] = "Bob"
|
||||
ws["B3"] = 25
|
||||
|
||||
# Add a second sheet
|
||||
ws2 = wb.create_sheet("Sheet2")
|
||||
ws2["A1"] = "Product"
|
||||
ws2["B1"] = "Price"
|
||||
ws2["A2"] = "Widget"
|
||||
ws2["B2"] = 9.99
|
||||
|
||||
wb.save(xlsx_file)
|
||||
wb.close()
|
||||
|
||||
result = extract_text(xlsx_file)
|
||||
assert result is not None
|
||||
assert "--- Sheet: Sheet1 ---" in result
|
||||
assert "--- Sheet: Sheet2 ---" in result
|
||||
assert "Alice" in result
|
||||
assert "Bob" in result
|
||||
assert "Widget" in result
|
||||
assert "9.99" in result
|
||||
|
||||
def test_extract_text_xlsx_empty_sheet(self, tmp_path: Path):
|
||||
"""Test extracting text from an .xlsx file with empty sheets."""
|
||||
from openpyxl import Workbook
|
||||
|
||||
xlsx_file = tmp_path / "empty.xlsx"
|
||||
wb = Workbook()
|
||||
# Clear the default sheet
|
||||
wb.remove(wb.active)
|
||||
# Add an empty sheet
|
||||
wb.create_sheet("EmptySheet")
|
||||
wb.save(xlsx_file)
|
||||
wb.close()
|
||||
|
||||
result = extract_text(xlsx_file)
|
||||
# Empty sheets should return empty string or header only
|
||||
assert result == "--- Sheet: EmptySheet ---" or result == ""
|
||||
|
||||
def test_extract_text_docx(self, tmp_path: Path):
|
||||
"""Test extracting text from a .docx file."""
|
||||
from docx import Document
|
||||
|
||||
docx_file = tmp_path / "test.docx"
|
||||
doc = Document()
|
||||
doc.add_heading("Test Document", 0)
|
||||
doc.add_paragraph("This is paragraph one.")
|
||||
doc.add_paragraph("This is paragraph two.")
|
||||
doc.save(docx_file)
|
||||
|
||||
result = extract_text(docx_file)
|
||||
assert result is not None
|
||||
assert "Test Document" in result
|
||||
assert "This is paragraph one." in result
|
||||
assert "This is paragraph two." in result
|
||||
|
||||
def test_extract_text_docx_empty(self, tmp_path: Path):
|
||||
"""Test extracting text from an empty .docx file."""
|
||||
from docx import Document
|
||||
|
||||
docx_file = tmp_path / "empty.docx"
|
||||
doc = Document()
|
||||
doc.save(docx_file)
|
||||
|
||||
result = extract_text(docx_file)
|
||||
assert result == ""
|
||||
|
||||
def test_extract_text_pptx(self, tmp_path: Path):
|
||||
"""Test extracting text from a .pptx file."""
|
||||
from pptx import Presentation
|
||||
|
||||
pptx_file = tmp_path / "test.pptx"
|
||||
prs = Presentation()
|
||||
|
||||
# Slide 1
|
||||
slide1 = prs.slides.add_slide(prs.slide_layouts[0])
|
||||
for shape in slide1.shapes:
|
||||
if hasattr(shape, "text"):
|
||||
shape.text = "First Slide Title"
|
||||
|
||||
# Slide 2
|
||||
slide2 = prs.slides.add_slide(prs.slide_layouts[5])
|
||||
left = top = width = height = 1000000
|
||||
textbox = slide2.shapes.add_textbox(left, top, width, height)
|
||||
text_frame = textbox.text_frame
|
||||
text_frame.text = "Bullet point content"
|
||||
|
||||
prs.save(pptx_file)
|
||||
|
||||
result = extract_text(pptx_file)
|
||||
assert result is not None
|
||||
assert "--- Slide 1 ---" in result
|
||||
assert "--- Slide 2 ---" in result
|
||||
# Text content may vary depending on PowerPoint layout defaults
|
||||
assert len(result) > 0
|
||||
|
||||
def test_extract_text_pdf_not_found(self, tmp_path: Path):
|
||||
"""Test that missing PDF files return error string."""
|
||||
missing_pdf = tmp_path / "nonexistent.pdf"
|
||||
|
||||
result = extract_text(missing_pdf)
|
||||
assert result is not None
|
||||
assert "[error: file not found:" in result
|
||||
|
||||
def test_extract_text_image_files(self, tmp_path: Path):
|
||||
"""Test that image files return placeholder text."""
|
||||
# Create a minimal PNG file (1x1 pixel)
|
||||
png_file = tmp_path / "test.png"
|
||||
# Minimal valid PNG: 8-byte signature + IHDR + IDAT + IEND
|
||||
png_data = (
|
||||
b"\x89PNG\r\n\x1a\n"
|
||||
b"\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
|
||||
b"\x08\x02\x00\x00\x00\x90wS\xde"
|
||||
b"\x00\x00\x00\x0cIDATx\x9cc\x00\x01\x00\x00\x05\x00\x01"
|
||||
b"\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
png_file.write_bytes(png_data)
|
||||
|
||||
result = extract_text(png_file)
|
||||
assert result is not None
|
||||
assert "[image:" in result
|
||||
assert "test.png" in result
|
||||
|
||||
|
||||
class TestIsTextExtension:
|
||||
"""Test the _is_text_extension helper."""
|
||||
|
||||
def test_text_extensions_return_true(self):
|
||||
"""Test that known text extensions return True."""
|
||||
assert _is_text_extension(".txt") is True
|
||||
assert _is_text_extension(".md") is True
|
||||
assert _is_text_extension(".csv") is True
|
||||
assert _is_text_extension(".json") is True
|
||||
assert _is_text_extension(".yaml") is True
|
||||
assert _is_text_extension(".yml") is True
|
||||
assert _is_text_extension(".xml") is True
|
||||
assert _is_text_extension(".html") is True
|
||||
assert _is_text_extension(".htm") is True
|
||||
|
||||
def test_non_text_extensions_return_false(self):
|
||||
"""Test that non-text extensions return False."""
|
||||
assert _is_text_extension(".pdf") is False
|
||||
assert _is_text_extension(".docx") is False
|
||||
assert _is_text_extension(".xlsx") is False
|
||||
assert _is_text_extension(".pptx") is False
|
||||
assert _is_text_extension(".png") is False
|
||||
assert _is_text_extension(".xyz") is False
|
||||
|
||||
def test_case_sensitivity(self):
|
||||
"""Test that _is_text_extension requires lowercase extension.
|
||||
|
||||
Note: The main extract_text function handles case-insensitivity by
|
||||
converting extensions to lowercase before calling _is_text_extension.
|
||||
"""
|
||||
# _is_text_extension itself is case-sensitive (lowercase only)
|
||||
assert _is_text_extension(".txt") is True
|
||||
assert _is_text_extension(".TXT") is False
|
||||
assert _is_text_extension(".pdf") is False
|
||||
+39
-10
@@ -194,6 +194,7 @@ async def test_successful_request_uses_fixed_api_session(aiohttp_client, mock_ag
|
||||
assert body["model"] == "test-model"
|
||||
mock_agent.process_direct.assert_called_once_with(
|
||||
content="hello",
|
||||
media=None,
|
||||
session_key=API_SESSION_KEY,
|
||||
channel="api",
|
||||
chat_id=API_CHAT_ID,
|
||||
@@ -205,7 +206,7 @@ async def test_successful_request_uses_fixed_api_session(aiohttp_client, mock_ag
|
||||
async def test_followup_requests_share_same_session_key(aiohttp_client) -> None:
|
||||
call_log: list[str] = []
|
||||
|
||||
async def fake_process(content, session_key="", channel="", chat_id=""):
|
||||
async def fake_process(content, session_key="", channel="", chat_id="", **kwargs):
|
||||
call_log.append(session_key)
|
||||
return f"reply to {content}"
|
||||
|
||||
@@ -236,7 +237,7 @@ async def test_followup_requests_share_same_session_key(aiohttp_client) -> None:
|
||||
async def test_fixed_session_requests_are_serialized(aiohttp_client) -> None:
|
||||
order: list[str] = []
|
||||
|
||||
async def slow_process(content, session_key="", channel="", chat_id=""):
|
||||
async def slow_process(content, session_key="", channel="", chat_id="", **kwargs):
|
||||
order.append(f"start:{content}")
|
||||
await asyncio.sleep(0.1)
|
||||
order.append(f"end:{content}")
|
||||
@@ -307,12 +308,12 @@ async def test_multimodal_content_extracts_text(aiohttp_client, mock_agent) -> N
|
||||
},
|
||||
)
|
||||
assert resp.status == 200
|
||||
mock_agent.process_direct.assert_called_once_with(
|
||||
content="describe this",
|
||||
session_key=API_SESSION_KEY,
|
||||
channel="api",
|
||||
chat_id=API_CHAT_ID,
|
||||
)
|
||||
call_kwargs = mock_agent.process_direct.call_args.kwargs
|
||||
assert call_kwargs["content"] == "describe this"
|
||||
assert call_kwargs["session_key"] == API_SESSION_KEY
|
||||
assert call_kwargs["channel"] == "api"
|
||||
assert call_kwargs["chat_id"] == API_CHAT_ID
|
||||
assert len(call_kwargs.get("media") or []) >= 0 # base64 images saved to disk
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
|
||||
@@ -320,7 +321,7 @@ async def test_multimodal_content_extracts_text(aiohttp_client, mock_agent) -> N
|
||||
async def test_empty_response_retry_then_success(aiohttp_client) -> None:
|
||||
call_count = 0
|
||||
|
||||
async def sometimes_empty(content, session_key="", channel="", chat_id=""):
|
||||
async def sometimes_empty(content, session_key="", channel="", chat_id="", **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
@@ -351,7 +352,7 @@ async def test_empty_response_falls_back(aiohttp_client) -> None:
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def always_empty(content, session_key="", channel="", chat_id=""):
|
||||
async def always_empty(content, session_key="", channel="", chat_id="", **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return ""
|
||||
@@ -371,3 +372,31 @@ async def test_empty_response_falls_back(aiohttp_client) -> None:
|
||||
body = await resp.json()
|
||||
assert body["choices"][0]["message"]["content"] == EMPTY_FINAL_RESPONSE_MESSAGE
|
||||
assert call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_direct_accepts_media() -> None:
|
||||
"""process_direct should forward media paths to _process_message."""
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
|
||||
loop = AgentLoop.__new__(AgentLoop)
|
||||
loop._connect_mcp = AsyncMock()
|
||||
|
||||
captured_msg = None
|
||||
|
||||
async def fake_process(msg, *, session_key="", on_progress=None, on_stream=None, on_stream_end=None):
|
||||
nonlocal captured_msg
|
||||
captured_msg = msg
|
||||
return None
|
||||
|
||||
loop._process_message = fake_process
|
||||
|
||||
await loop.process_direct(
|
||||
content="analyze this",
|
||||
media=["/tmp/image.png", "/tmp/report.pdf"],
|
||||
session_key="test:1",
|
||||
)
|
||||
|
||||
assert captured_msg is not None
|
||||
assert captured_msg.media == ["/tmp/image.png", "/tmp/report.pdf"]
|
||||
assert captured_msg.content == "analyze this"
|
||||
|
||||
Reference in New Issue
Block a user