zk-data-agent/tests/test_evaluation_api.py

from __future__ import annotations

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

from fastapi.testclient import TestClient

from backend.api.server import AgentState, create_app


def _build_state(root: Path) -> AgentState:
    skill_dir = root / 'skills' / 'label-master'
    manifest_dir = skill_dir / 'knowledge' / '索引'
    manifest_dir.mkdir(parents=True)
    (skill_dir / 'SKILL.md').write_text(
        (
            '---\n'
            'name: label-master\n'
            'description: Test label master.\n'
            'allowed_tools: read_file\n'
            '---\n'
            'Use the label catalog.\n'
        ),
        encoding='utf-8',
    )
    (manifest_dir / 'label_manifest.json').write_text(
        json.dumps(
            {
                'counts': {'labels': 1},
                'labels': [
                    {
                        'name': '地图导航',
                        'domain': '地图',
                        'path': 'knowledge/标签/地图导航.md',
                    }
                ],
                'agent_tags': ['地图导航'],
                'functions': [],
                'intents': [],
            },
            ensure_ascii=False,
        ),
        encoding='utf-8',
    )
    return AgentState(
        cwd=root,
        model='test-model',
        base_url='http://127.0.0.1:8000/v1',
        api_key='local-token',
        timeout_seconds=30,
        allow_shell=False,
        allow_write=False,
        session_directory=root / '.port_sessions' / 'agent',
    )


class EvaluationApiTests(unittest.TestCase):
    def test_single_analysis_endpoint(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            state = _build_state(root)
            with (
                TestClient(create_app(state)) as client,
                patch.object(
                    state.evaluation_runtime,
                    'analyze_case',
                    return_value={
                        'query': '导航去公司',
                        'prediction': '地图导航',
                        'status': 'completed',
                    },
                ) as analyze,
            ):
                response = client.post(
                    '/api/evaluations/analyze',
                    json={
                        'account_id': 'alice',
                        'query': '导航去公司',
                        'skill_name': 'label-master',
                        'model': 'test-model',
                    },
                )
                self.assertEqual(response.status_code, 200)
                self.assertEqual(response.json()['prediction'], '地图导航')
                analyze.assert_called_once()

    def test_single_analysis_history_endpoints(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            state = _build_state(root)
            record = {
                'id': 'single_1',
                'query': '导航去公司',
                'prediction': '地图导航',
                'status': 'completed',
            }
            with (
                TestClient(create_app(state)) as client,
                patch.object(
                    state.evaluation_runtime,
                    'list_single_analyses',
                    return_value=[record],
                ) as list_history,
                patch.object(
                    state.evaluation_runtime,
                    'get_single_analysis',
                    return_value=record,
                ) as get_history,
            ):
                response = client.get(
                    '/api/evaluations/analyses',
                    params={'account_id': 'alice'},
                )
                self.assertEqual(response.status_code, 200)
                self.assertEqual(response.json(), [record])
                list_history.assert_called_once_with('alice')

                detail = client.get(
                    '/api/evaluations/analyses/single_1',
                    params={'account_id': 'alice'},
                )
                self.assertEqual(detail.status_code, 200)
                self.assertEqual(detail.json(), record)
                get_history.assert_called_once_with('single_1', 'alice')

    def test_dataset_experiment_and_export_flow(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            state = _build_state(root)
            with TestClient(create_app(state)) as client:
                metadata = client.get(
                    '/api/evaluations/metadata',
                    params={'account_id': 'alice'},
                )
                self.assertEqual(metadata.status_code, 200)
                self.assertEqual(metadata.json()['default_skill'], 'label-master')

                encoded = base64.b64encode(
                    'query,label\n导航去公司,地图导航\n'.encode()
                ).decode()
                dataset_response = client.post(
                    '/api/evaluations/datasets',
                    json={
                        'account_id': 'alice',
                        'name': 'routing',
                        'filename': 'routing.csv',
                        'content_base64': encoded,
                    },
                )
                self.assertEqual(dataset_response.status_code, 200)
                dataset = dataset_response.json()
                self.assertEqual(dataset['row_count'], 1)

                experiment_response = client.post(
                    '/api/evaluations/experiments',
                    json={
                        'account_id': 'alice',
                        'dataset_id': dataset['id'],
                        'name': 'routing test',
                        'skill_name': 'label-master',
                        'concurrency': 1,
                    },
                )
                self.assertEqual(experiment_response.status_code, 200)
                experiment = experiment_response.json()
                self.assertEqual(experiment['status'], 'draft')
                self.assertEqual(experiment['cases'][0]['query'], '导航去公司')

                export = client.get(
                    f"/api/evaluations/experiments/{experiment['id']}/export",
                    params={'account_id': 'alice'},
                )
                self.assertEqual(export.status_code, 200)
                self.assertIn('text/csv', export.headers['content-type'])
                self.assertIn('导航去公司', export.content.decode('utf-8-sig'))


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