"""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}]