diff options
Diffstat (limited to 'tests')
| -rw-r--r-- | tests/test_clear_messages.py | 63 | ||||
| -rw-r--r-- | tests/test_llm.py | 83 |
2 files changed, 146 insertions, 0 deletions
diff --git a/tests/test_clear_messages.py b/tests/test_clear_messages.py new file mode 100644 index 0000000..91a2556 --- /dev/null +++ b/tests/test_clear_messages.py @@ -0,0 +1,63 @@ +"""Tests for clearing session messages and files.""" + +from __future__ import annotations + +import json +import tempfile +from pathlib import Path + +from ollama_web.sessions import SessionStore, is_valid_id + + +def make_store() -> SessionStore: + tmp = tempfile.mkdtemp() + return SessionStore(data_dir=tmp) + + +def test_clear_messages_removes_history_and_files(): + store = make_store() + session = store.create(title="test") + session_id = session["id"] + assert is_valid_id(session_id) + + chat_file = store.add_file( + session_id=session_id, + name="notes.txt", + data=b"hello", + mime="text/plain", + text="hello", + ) + store.add_message( + session_id=session_id, + role="user", + content="hi", + attachments=[chat_file.id], + ) + store.add_message(session_id=session_id, role="assistant", content="hello") + + files_dir = store._files_dir(session_id) + stored_path = store.data_dir / chat_file.path + assert stored_path.exists() + assert (files_dir / f"{chat_file.id}.txt").exists() + + cleared = store.clear_messages(session_id) + assert cleared is not None + assert cleared["messages"] == [] + assert cleared["files"] == [] + assert not stored_path.exists() + assert files_dir.exists() + + reloaded = store.get(session_id) + assert reloaded is not None + assert reloaded["messages"] == [] + assert reloaded["files"] == [] + + +def test_clear_messages_invalid_session_id(): + store = make_store() + assert store.clear_messages("not-a-valid-id") is None + + +def test_clear_messages_missing_session(): + store = make_store() + assert store.clear_messages("0" * 32) is None
\ No newline at end of file diff --git a/tests/test_llm.py b/tests/test_llm.py new file mode 100644 index 0000000..cbb9910 --- /dev/null +++ b/tests/test_llm.py @@ -0,0 +1,83 @@ +"""Tests for llm capability detection and feature fallbacks.""" + +from __future__ import annotations + +import base64 + +import pytest +from ollama import Message +from ollama._types import Image as OllamaImage +from ollama._types import ResponseError + +from ollama_web.llm import ( + _embed_images_in_content, + _is_image_message, + chat_with_tools, + stream_chat_with_tools, +) +from ollama_web.tools.registry import ToolRegistry + + +def _image_message() -> Message: + return Message( + role="user", + content="hello", + images=[OllamaImage(value=b"fake-image-bytes")], + ) + + +def test_is_image_message_detects_images() -> None: + assert _is_image_message(_image_message()) is True + + +def test_is_image_message_without_images() -> None: + assert _is_image_message(Message(role="user", content="hello")) is False + + +def test_embed_images_in_content_adds_base64_markdown() -> None: + embedded = _embed_images_in_content(_image_message()) + encoded = base64.b64encode(b"fake-image-bytes").decode("ascii") + assert f"" in embedded.content + assert not getattr(embedded, "images", None) + + +def test_chat_with_tools_skips_tools_for_non_tool_model() -> None: + """Tools are not passed when the model does not advertise tool support.""" + messages = [Message(role="user", content="hi")] + # Use an empty registry so any tool call would fail if requested. + registry = ToolRegistry() + + # Without tool capability, the iterator should finish immediately with no + # actual network request. We only verify it accepts the call and does not + # raise due to tool-related setup. + gen = chat_with_tools( + messages, + model="nonexistent-model", + capabilities=set(), + registry=registry, + ) + events = list(gen) + # We expect either an error (model not found) or a done event. + assert any(e.get("type") in {"error", "done"} for e in events) + + +def test_chat_with_tools_forces_think_none_for_non_thinking_model() -> None: + messages = [Message(role="user", content="hi")] + gen = chat_with_tools( + messages, + model="nonexistent-model", + capabilities=set(), + ) + events = list(gen) + assert any(e.get("type") in {"error", "done"} for e in events) + + +def test_stream_chat_with_tools_accepts_capability_set() -> None: + messages = [Message(role="user", content="hi")] + gen = stream_chat_with_tools( + messages, + model="nonexistent-model", + capabilities={"tools", "thinking"}, + ) + with pytest.raises(ResponseError): + list(gen) |