486 lines
16 KiB
Python
486 lines
16 KiB
Python
"""Tests for BitableTool (U4).
|
|
|
|
Tests the full HTTP flow: BitableTool → bitable REST API → BitableService.
|
|
Uses ``httpx.AsyncClient`` + ``ASGITransport`` so the tool's HTTP calls
|
|
and the bitable DB share one event loop.
|
|
|
|
Covers:
|
|
- KTD11: X-Internal-Token auth (valid token accepted, invalid rejected)
|
|
- Batch chunking: 1200 records → 3 HTTP requests (500+500+200)
|
|
- Resume from partial failure
|
|
- Three ingestion types: Excel, database, API collector
|
|
- create_table, upsert_records, query_records
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
|
|
import httpx
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from httpx import ASGITransport
|
|
|
|
from agentkit.bitable.service import BitableService
|
|
from agentkit.server.routes import bitable as bitable_routes
|
|
from agentkit.server.routes.bitable import require_bitable_auth
|
|
from agentkit.tools.bitable_tool import BATCH_SIZE, BitableTool
|
|
|
|
pytestmark = pytest.mark.postgres
|
|
|
|
TEST_TOKEN = "test-internal-token-abc123"
|
|
TEST_USER = {"user_id": "test-user", "username": "tester", "role": "member"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def app(bitable_service: BitableService) -> FastAPI:
|
|
"""Test app with bitable_service + internal token on app.state."""
|
|
app = FastAPI()
|
|
app.state.bitable_service = bitable_service
|
|
app.state.bitable_internal_token = TEST_TOKEN
|
|
app.include_router(bitable_routes.router, prefix="/api/v1")
|
|
# Override auth so JWT path also works (for non-internal-token tests)
|
|
app.dependency_overrides[require_bitable_auth] = lambda: TEST_USER
|
|
return app
|
|
|
|
|
|
@pytest.fixture
|
|
def app_no_override(bitable_service: BitableService) -> FastAPI:
|
|
"""App without auth override — tests real X-Internal-Token path."""
|
|
app = FastAPI()
|
|
app.state.bitable_service = bitable_service
|
|
app.state.bitable_internal_token = TEST_TOKEN
|
|
app.include_router(bitable_routes.router, prefix="/api/v1")
|
|
return app
|
|
|
|
|
|
@pytest.fixture
|
|
def app_no_token(bitable_service: BitableService) -> FastAPI:
|
|
"""App without internal token configured."""
|
|
app = FastAPI()
|
|
app.state.bitable_service = bitable_service
|
|
app.include_router(bitable_routes.router, prefix="/api/v1")
|
|
app.dependency_overrides[require_bitable_auth] = lambda: TEST_USER
|
|
return app
|
|
|
|
|
|
def _make_client(app: FastAPI, token: str | None = None) -> httpx.AsyncClient:
|
|
"""Create an httpx AsyncClient backed by ASGITransport.
|
|
|
|
If token is provided, the X-Internal-Token header is set as default
|
|
on the client — mirroring how BitableTool._get_client configures it.
|
|
"""
|
|
base = "http://test/api/v1/bitable"
|
|
transport = ASGITransport(app=app)
|
|
headers: dict[str, str] = {}
|
|
if token:
|
|
headers["X-Internal-Token"] = token
|
|
return httpx.AsyncClient(transport=transport, base_url=base, headers=headers)
|
|
|
|
|
|
@pytest.fixture
|
|
async def tool(app: FastAPI) -> BitableTool:
|
|
"""BitableTool pointing at the test app via ASGITransport.
|
|
|
|
ponytail: We patch _client to use ASGITransport instead of real
|
|
HTTP — this shares the event loop with the async DB fixtures.
|
|
"""
|
|
client = _make_client(app, token=TEST_TOKEN)
|
|
t = BitableTool(base_url="http://test/api/v1/bitable", internal_token=TEST_TOKEN)
|
|
t._client = client
|
|
yield t
|
|
await client.aclose()
|
|
|
|
|
|
@pytest.fixture
|
|
async def tool_no_token(app_no_token: FastAPI) -> BitableTool:
|
|
"""BitableTool without internal token."""
|
|
client = _make_client(app_no_token, token=None)
|
|
t = BitableTool(base_url="http://test/api/v1/bitable", internal_token=None)
|
|
t._client = client
|
|
yield t
|
|
await client.aclose()
|
|
|
|
|
|
@pytest.fixture
|
|
async def tool_real_auth(app_no_override: FastAPI) -> BitableTool:
|
|
"""BitableTool that sends real X-Internal-Token header (no auth override)."""
|
|
client = _make_client(app_no_override, token=TEST_TOKEN)
|
|
t = BitableTool(base_url="http://test/api/v1/bitable", internal_token=TEST_TOKEN)
|
|
t._client = client
|
|
yield t
|
|
await client.aclose()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_xlsx(sheets: dict[str, list[list]]) -> bytes:
|
|
"""Create an in-memory .xlsx file."""
|
|
from openpyxl import Workbook
|
|
|
|
wb = Workbook()
|
|
wb.remove(wb.active)
|
|
for name, rows in sheets.items():
|
|
ws = wb.create_sheet(title=name)
|
|
for row in rows:
|
|
ws.append(row)
|
|
buf = io.BytesIO()
|
|
wb.save(buf)
|
|
return buf.getvalue()
|
|
|
|
|
|
async def _setup_table_with_pk(tool: BitableTool, name: str = "T") -> tuple[str, str, str]:
|
|
"""Create a table with a text PK field and a number data field.
|
|
|
|
Returns (table_id, pk_field_id, data_field_id).
|
|
"""
|
|
result = await tool.execute(action="create_table", table_name=name)
|
|
assert result["success"], result
|
|
table_id = result["table"]["id"]
|
|
|
|
client = await tool._get_client()
|
|
|
|
# Create PK field
|
|
resp = await client.post(
|
|
f"/tables/{table_id}/fields",
|
|
json={"name": "id", "field_type": "text", "owner": "agent"},
|
|
)
|
|
resp.raise_for_status()
|
|
pk_field_id = resp.json()["field"]["id"]
|
|
|
|
# Create data field
|
|
resp = await client.post(
|
|
f"/tables/{table_id}/fields",
|
|
json={"name": "val", "field_type": "number", "owner": "agent"},
|
|
)
|
|
resp.raise_for_status()
|
|
data_field_id = resp.json()["field"]["id"]
|
|
|
|
# Set PK
|
|
resp = await client.patch(f"/tables/{table_id}", json={"primary_key_field_id": pk_field_id})
|
|
resp.raise_for_status()
|
|
|
|
return table_id, pk_field_id, data_field_id
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# create_table
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_create_table(tool: BitableTool) -> None:
|
|
"""create_table action creates a bitable table via HTTP."""
|
|
result = await tool.execute(action="create_table", table_name="MyTable")
|
|
assert result["success"] is True
|
|
assert result["table"]["name"] == "MyTable"
|
|
|
|
|
|
async def test_create_table_missing_name(tool: BitableTool) -> None:
|
|
"""Missing table_name → error."""
|
|
result = await tool.execute(action="create_table")
|
|
assert result["success"] is False
|
|
assert "table_name" in result["error"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# KTD11: Internal token auth
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_internal_token_accepted(tool_real_auth: BitableTool) -> None:
|
|
"""Valid X-Internal-Token → request succeeds (no JWT needed)."""
|
|
result = await tool_real_auth.execute(action="create_table", table_name="Authed")
|
|
assert result["success"] is True
|
|
|
|
|
|
async def test_invalid_token_rejected(app_no_override: FastAPI) -> None:
|
|
"""Wrong X-Internal-Token → 401."""
|
|
transport = ASGITransport(app=app_no_override)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post(
|
|
"/api/v1/bitable/tables",
|
|
json={"name": "X"},
|
|
headers={"X-Internal-Token": "wrong-token"},
|
|
)
|
|
assert resp.status_code == 401
|
|
|
|
|
|
async def test_no_auth_rejected(app_no_override: FastAPI) -> None:
|
|
"""No auth at all → 401."""
|
|
transport = ASGITransport(app=app_no_override)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
|
resp = await client.post("/api/v1/bitable/tables", json={"name": "X"})
|
|
assert resp.status_code == 401
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Batch chunking (BATCH_SIZE=500)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_batch_upsert_1200_records(tool: BitableTool) -> None:
|
|
"""1200 records → 3 batches (500+500+200), all succeed."""
|
|
table_id, pk_fid, data_fid = await _setup_table_with_pk(tool)
|
|
|
|
records = [{pk_fid: f"r{i}", data_fid: i * 10} for i in range(1200)]
|
|
result = await tool.execute(
|
|
action="upsert_records",
|
|
table_id=table_id,
|
|
records=records,
|
|
primary_key_field_id=pk_fid,
|
|
)
|
|
assert result["success"] is True
|
|
assert result["successful_count"] == 1200
|
|
assert result["total"] == 1200
|
|
assert "errors" not in result
|
|
|
|
|
|
async def test_batch_size_is_500() -> None:
|
|
"""Verify BATCH_SIZE constant is 500."""
|
|
assert BATCH_SIZE == 500
|
|
|
|
|
|
async def test_resume_from_partial_failure(tool: BitableTool) -> None:
|
|
"""resume_from skips already-successful records."""
|
|
table_id, pk_fid, data_fid = await _setup_table_with_pk(tool)
|
|
|
|
# First, insert 500 records successfully
|
|
batch1 = [{pk_fid: f"r{i}", data_fid: i} for i in range(500)]
|
|
result1 = await tool.execute(
|
|
action="upsert_records",
|
|
table_id=table_id,
|
|
records=batch1,
|
|
primary_key_field_id=pk_fid,
|
|
)
|
|
assert result1["successful_count"] == 500
|
|
|
|
# Now resume from 500 with the remaining 700
|
|
all_records = [{pk_fid: f"r{i}", data_fid: i} for i in range(1200)]
|
|
remaining = all_records[500:]
|
|
result2 = await tool.execute(
|
|
action="upsert_records",
|
|
table_id=table_id,
|
|
records=remaining,
|
|
primary_key_field_id=pk_fid,
|
|
resume_from=0, # remaining is already sliced
|
|
)
|
|
assert result2["successful_count"] == 700
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# query_records
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_query_records(tool: BitableTool) -> None:
|
|
"""query_records returns records from the table."""
|
|
table_id, pk_fid, data_fid = await _setup_table_with_pk(tool)
|
|
|
|
# Insert some records
|
|
await tool.execute(
|
|
action="upsert_records",
|
|
table_id=table_id,
|
|
records=[{pk_fid: "a", data_fid: 1}, {pk_fid: "b", data_fid: 2}],
|
|
primary_key_field_id=pk_fid,
|
|
)
|
|
|
|
# Query
|
|
result = await tool.execute(action="query_records", table_id=table_id)
|
|
assert result["success"] is True
|
|
assert len(result["records"]) == 2
|
|
|
|
|
|
async def test_query_records_with_limit(tool: BitableTool) -> None:
|
|
"""query_records with limit returns fewer records."""
|
|
table_id, pk_fid, data_fid = await _setup_table_with_pk(tool)
|
|
|
|
await tool.execute(
|
|
action="upsert_records",
|
|
table_id=table_id,
|
|
records=[{pk_fid: f"r{i}", data_fid: i} for i in range(10)],
|
|
primary_key_field_id=pk_fid,
|
|
)
|
|
|
|
result = await tool.execute(action="query_records", table_id=table_id, limit=5)
|
|
assert result["success"] is True
|
|
assert len(result["records"]) == 5
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# import_excel
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_import_excel_file(tool: BitableTool, tmp_path) -> None:
|
|
"""import_excel from file path → creates table + fields + records."""
|
|
xlsx_bytes = _make_xlsx({"Products": [["name", "price"], ["Widget", 9.99], ["Gadget", 19.99]]})
|
|
file_path = tmp_path / "test.xlsx"
|
|
file_path.write_bytes(xlsx_bytes)
|
|
|
|
result = await tool.execute(action="import_excel", file_path=str(file_path))
|
|
assert result["success"] is True
|
|
sheet_result = result["sheets"][0]
|
|
assert sheet_result["record_count"] == 2
|
|
assert sheet_result["field_count"] == 2
|
|
|
|
# Verify data was actually written
|
|
table_id = sheet_result["table_id"]
|
|
query = await tool.execute(action="query_records", table_id=table_id)
|
|
assert len(query["records"]) == 2
|
|
|
|
|
|
async def test_import_excel_empty_sheet(tool: BitableTool, tmp_path) -> None:
|
|
"""Excel with only headers (no data rows) → table created, 0 records."""
|
|
xlsx_bytes = _make_xlsx({"Empty": [["col1", "col2"]]})
|
|
file_path = tmp_path / "empty.xlsx"
|
|
file_path.write_bytes(xlsx_bytes)
|
|
|
|
result = await tool.execute(action="import_excel", file_path=str(file_path))
|
|
assert result["success"] is True
|
|
assert result["sheets"][0]["record_count"] == 0
|
|
assert result["sheets"][0]["field_count"] == 2
|
|
|
|
|
|
async def test_import_excel_missing_path(tool: BitableTool) -> None:
|
|
"""No file_path or file_url → error."""
|
|
result = await tool.execute(action="import_excel")
|
|
assert result["success"] is False
|
|
assert "file_path" in result["error"] or "file_url" in result["error"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# collect_api
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_collect_api(tool: BitableTool) -> None:
|
|
"""collect_api transforms records via field_mapping and upserts."""
|
|
table_id, pk_fid, data_fid = await _setup_table_with_pk(tool)
|
|
|
|
result = await tool.execute(
|
|
action="collect_api",
|
|
table_id=table_id,
|
|
records=[
|
|
{"user_id": "u1", "score": 100},
|
|
{"user_id": "u2", "score": 200},
|
|
],
|
|
field_mapping={"user_id": pk_fid, "score": data_fid},
|
|
primary_key_field_id=pk_fid,
|
|
)
|
|
assert result["success"] is True
|
|
assert result["successful_count"] == 2
|
|
|
|
# Verify
|
|
query = await tool.execute(action="query_records", table_id=table_id)
|
|
assert len(query["records"]) == 2
|
|
|
|
|
|
async def test_collect_api_missing_fields(tool: BitableTool) -> None:
|
|
"""Missing required fields → error."""
|
|
result = await tool.execute(action="collect_api", records=[])
|
|
assert result["success"] is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Error handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_unknown_action(tool: BitableTool) -> None:
|
|
"""Unknown action → error."""
|
|
result = await tool.execute(action="bogus")
|
|
assert result["success"] is False
|
|
assert "Unknown action" in result["error"]
|
|
|
|
|
|
async def test_query_nonexistent_table(tool: BitableTool) -> None:
|
|
"""Querying a non-existent table → error."""
|
|
result = await tool.execute(action="query_records", table_id="nonexistent-id")
|
|
assert result["success"] is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Database ingestion (type mapping only — no real external DB needed)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_db_type_mapping_integer() -> None:
|
|
"""Integer type → 'number'."""
|
|
from sqlalchemy import Integer
|
|
|
|
from agentkit.bitable.ingestion.database import infer_field_type
|
|
|
|
assert infer_field_type(Integer()) == "number"
|
|
assert infer_field_type(Integer) == "number"
|
|
|
|
|
|
def test_db_type_mapping_varchar() -> None:
|
|
"""String type → 'text'."""
|
|
from sqlalchemy import String
|
|
|
|
from agentkit.bitable.ingestion.database import infer_field_type
|
|
|
|
assert infer_field_type(String(255)) == "text"
|
|
|
|
|
|
def test_db_type_mapping_datetime() -> None:
|
|
"""DateTime type → 'date'."""
|
|
from sqlalchemy import DateTime
|
|
|
|
from agentkit.bitable.ingestion.database import infer_field_type
|
|
|
|
assert infer_field_type(DateTime()) == "date"
|
|
|
|
|
|
def test_db_type_mapping_unknown_fallback() -> None:
|
|
"""Unknown type → 'text' (safe fallback)."""
|
|
from agentkit.bitable.ingestion.database import infer_field_type
|
|
|
|
class CustomType:
|
|
pass
|
|
|
|
assert infer_field_type(CustomType()) == "text"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# API collector transform
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_transform_records_basic() -> None:
|
|
"""transform_records maps source keys to field IDs."""
|
|
from agentkit.bitable.ingestion.api_collector import transform_records
|
|
|
|
result = transform_records(
|
|
records=[{"name": "Alice", "age": 30, "extra": "dropped"}],
|
|
field_mapping={"name": "fld_abc", "age": "fld_def"},
|
|
)
|
|
assert result == [{"fld_abc": "Alice", "fld_def": 30}]
|
|
|
|
|
|
def test_transform_records_empty() -> None:
|
|
"""Empty records → empty result."""
|
|
from agentkit.bitable.ingestion.api_collector import transform_records
|
|
|
|
assert transform_records([], {"a": "b"}) == []
|
|
assert transform_records([{"a": 1}], {}) == []
|
|
|
|
|
|
def test_transform_records_missing_keys() -> None:
|
|
"""Source keys not in mapping are silently dropped."""
|
|
from agentkit.bitable.ingestion.api_collector import transform_records
|
|
|
|
result = transform_records(
|
|
records=[{"a": 1, "b": 2}],
|
|
field_mapping={"a": "fld_a"}, # b is not mapped
|
|
)
|
|
assert result == [{"fld_a": 1}]
|