zk-data-agent/tests/test_feishu_integration.py

from __future__ import annotations

import tempfile
import unittest
import zipfile
import json
from pathlib import Path
from unittest.mock import patch

import openpyxl
from fastapi.testclient import TestClient

from backend.api import server as gui_server
from backend.api.server import AgentState, create_app


def _build_client(tmp: Path) -> tuple[TestClient, AgentState]:
    state = AgentState(
        cwd=tmp,
        model='test-model',
        base_url='http://127.0.0.1:8000/v1',
        api_key='local-token',
        allow_shell=False,
        allow_write=False,
        session_directory=tmp / 'sessions',
    )
    return TestClient(create_app(state)), state


class FeishuIntegrationTests(unittest.TestCase):
    def test_convert_xlsx_and_docx_to_markdown(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            root = Path(tmp_dir)
            xlsx_path = root / '样例.xlsx'
            docx_path = root / '说明.docx'
            _write_xlsx(xlsx_path)
            _write_docx(docx_path)

            xlsx_markdown = gui_server._convert_file_to_feishu_markdown(
                xlsx_path,
                title='表格样例',
            )
            docx_markdown = gui_server._convert_file_to_feishu_markdown(
                docx_path,
                title='文档样例',
            )

        self.assertIn('| query | label |', xlsx_markdown)
        self.assertIn('| 怎么去公司 | Agent(tag="地图导航") |', xlsx_markdown)
        self.assertIn('产品定义说明', docx_markdown)
        self.assertIn('| 功能点 | 示例 |', docx_markdown)

    def test_create_online_doc_requires_feishu_login(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            client, state = _build_client(Path(tmp_dir))
            file_path = _write_account_file(state, 'alice', 'draft.md', '# hello')

            with patch.object(
                gui_server,
                '_feishu_status_payload',
                return_value={'logged_in': False, 'status': 'not_logged_in'},
            ):
                response = client.post(
                    '/api/files/online-doc',
                    json={'account_id': 'alice', 'path': str(file_path)},
                )

        self.assertEqual(response.status_code, 409)
        self.assertEqual(response.json()['detail']['code'], 'feishu_not_logged_in')

    def test_create_online_doc_calls_feishu_mcp(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            client, state = _build_client(Path(tmp_dir))
            file_path = _write_account_file(state, 'alice', 'draft.md', '# hello')

            with patch.object(
                gui_server,
                '_feishu_status_payload',
                return_value={'logged_in': True, 'status': 'logged_in'},
            ), patch.object(
                gui_server.MCPRuntime,
                'call_tool',
                return_value=(
                    '{"url":"https://mi.feishu.cn/docx/example"}',
                    {'server_name': 'feishu-mcp-pro', 'tool_name': 'doc_create'},
                ),
            ) as call_tool:
                response = client.post(
                    '/api/files/online-doc',
                    json={'account_id': 'alice', 'path': str(file_path)},
                )

            self.assertEqual(response.status_code, 200)
            payload = response.json()
            self.assertEqual(payload['url'], 'https://mi.feishu.cn/docx/example')
            call_tool.assert_called_once()
            self.assertEqual(call_tool.call_args.args[0], 'doc_create')
            map_path = (
                state.account_paths('alice')['base']
                / 'integrations'
                / 'feishu'
                / 'online-docs.json'
            )
            online_docs = json.loads(map_path.read_text(encoding='utf-8'))
            self.assertEqual(
                online_docs['files'][str(file_path)]['url'],
                'https://mi.feishu.cn/docx/example',
            )

    def test_extract_first_url_ignores_wrapping_quotes(self) -> None:
        url = gui_server._extract_first_url(
            '{"url":"https://mi.feishu.cn/docx/CZIldpM3QofbbXxxAq1cEininTd"}'
        )

        self.assertEqual(url, 'https://mi.feishu.cn/docx/CZIldpM3QofbbXxxAq1cEininTd')

    def test_create_online_doc_rejects_jsonl_for_now(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            client, state = _build_client(Path(tmp_dir))
            file_path = _write_account_file(state, 'alice', 'records.jsonl', '{}\n')

            response = client.post(
                '/api/files/online-doc',
                json={'account_id': 'alice', 'path': str(file_path)},
            )

        self.assertEqual(response.status_code, 400)
        self.assertIn('json/jsonl', response.json()['detail'])


def _write_account_file(
    state: AgentState,
    account_id: str,
    name: str,
    content: str,
) -> Path:
    path = state.account_paths(account_id)['sessions'] / 's1' / 'output' / name
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding='utf-8')
    return path


def _write_xlsx(path: Path) -> None:
    workbook = openpyxl.Workbook()
    sheet = workbook.active
    sheet.title = 'Sheet1'
    sheet.append(['query', 'label'])
    sheet.append(['怎么去公司', 'Agent(tag="地图导航")'])
    workbook.save(path)


def _write_docx(path: Path) -> None:
    document_xml = '''<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
  <w:body>
    <w:p><w:r><w:t>产品定义说明</w:t></w:r></w:p>
    <w:tbl>
      <w:tr>
        <w:tc><w:p><w:r><w:t>功能点</w:t></w:r></w:p></w:tc>
        <w:tc><w:p><w:r><w:t>示例</w:t></w:r></w:p></w:tc>
      </w:tr>
      <w:tr>
        <w:tc><w:p><w:r><w:t>导航</w:t></w:r></w:p></w:tc>
        <w:tc><w:p><w:r><w:t>导航到公司</w:t></w:r></w:p></w:tc>
      </w:tr>
    </w:tbl>
  </w:body>
</w:document>
'''
    with zipfile.ZipFile(path, 'w') as archive:
        archive.writestr('word/document.xml', document_xml)


if __name__ == '__main__':
    unittest.main()