mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-28 14:07:55 +00:00
Elliot SluskyandGitHub
9db21d37ef
Fix the Ruff formatter check on main by formatting tests changed in #639 and #640 with the repository's pinned Ruff 0.15.1. No behavior change.
124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
"""Tests for KnowledgeSQLTool."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from openjarvis.connectors.store import KnowledgeStore
|
|
from openjarvis.core.registry import ToolRegistry
|
|
|
|
|
|
@pytest.fixture()
|
|
def store(tmp_path: Path) -> KnowledgeStore:
|
|
ks = KnowledgeStore(str(tmp_path / "test.db"))
|
|
ks.store("Hello from Alice", source="imessage", author="Alice", doc_type="message")
|
|
ks.store(
|
|
"Hello from Alice again", source="imessage", author="Alice", doc_type="message"
|
|
)
|
|
ks.store("Meeting notes Q1", source="granola", author="Bob", doc_type="document")
|
|
ks.store("Email about Spain trip", source="gmail", author="Carol", doc_type="email")
|
|
return ks
|
|
|
|
|
|
def test_select_count(store: KnowledgeStore) -> None:
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(query="SELECT COUNT(*) as total FROM knowledge_chunks")
|
|
assert result.success
|
|
assert "4" in result.content
|
|
|
|
|
|
def test_group_by_author(store: KnowledgeStore) -> None:
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(
|
|
query=(
|
|
"SELECT author, COUNT(*) as n "
|
|
"FROM knowledge_chunks "
|
|
"GROUP BY author ORDER BY n DESC"
|
|
)
|
|
)
|
|
assert result.success
|
|
assert "Alice" in result.content
|
|
assert "2" in result.content
|
|
|
|
|
|
def test_rejects_non_select(store: KnowledgeStore) -> None:
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(query="DELETE FROM knowledge_chunks")
|
|
assert not result.success
|
|
assert "read-only" in result.content.lower() or "SELECT" in result.content
|
|
|
|
|
|
def test_rejects_drop(store: KnowledgeStore) -> None:
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(query="DROP TABLE knowledge_chunks")
|
|
assert not result.success
|
|
|
|
|
|
def test_allows_select_with_keyword_substring(store: KnowledgeStore) -> None:
|
|
"""A read-only SELECT must not be rejected because a column/alias merely
|
|
contains a write keyword as a substring (e.g. 'created' -> CREATE)."""
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(query="SELECT author AS created_author FROM knowledge_chunks")
|
|
assert result.success, result.content
|
|
assert "Alice" in result.content
|
|
|
|
|
|
def test_allows_keyword_inside_string_literal(store: KnowledgeStore) -> None:
|
|
"""A write keyword appearing only inside a string literal must not be
|
|
treated as a forbidden statement."""
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(
|
|
query="SELECT content FROM knowledge_chunks WHERE content LIKE '%delete%'"
|
|
)
|
|
assert result.success, result.content
|
|
|
|
|
|
def test_rejects_multi_statement(store: KnowledgeStore) -> None:
|
|
"""Multi-statement strings fail with a ToolResult, not an exception."""
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(query="SELECT 1; VACUUM")
|
|
assert not result.success
|
|
assert "error" in result.content.lower()
|
|
|
|
|
|
def test_handles_bad_sql(store: KnowledgeStore) -> None:
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(query="SELECT * FROM nonexistent_table")
|
|
assert not result.success
|
|
|
|
|
|
def test_filter_by_source(store: KnowledgeStore) -> None:
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
tool = KnowledgeSQLTool(store=store)
|
|
result = tool.execute(
|
|
query="SELECT title, author FROM knowledge_chunks WHERE source = 'gmail'"
|
|
)
|
|
assert result.success
|
|
assert "Carol" in result.content
|
|
|
|
|
|
def test_registered() -> None:
|
|
from openjarvis.tools.knowledge_sql import KnowledgeSQLTool
|
|
|
|
ToolRegistry.register_value("knowledge_sql", KnowledgeSQLTool)
|
|
assert ToolRegistry.contains("knowledge_sql")
|