sql_qa_pair/tests/test_generate_endpoint.py

import asyncio
import unittest
from pathlib import Path

from app.endpoints.generate_qa_pair import generate_qa_pair
from app.endpoints.schema import QARequest


class GenerateQAPairEndpointTest(unittest.TestCase):
    def test_get_dashboard_info_operator(self):
        from app.operators.get_dashboard_info import get_dashboard_info

        info = get_dashboard_info({
            "dashboard_id": "DB001",
            "card_ids": ["C0001"],
            "bbk": "BBK002",
        })

        self.assertEqual(info.dashboard_id, "DB001")
        self.assertGreaterEqual(len(info.cards), 1)
        self.assertEqual(info.cards[0].card_id, "C0001")

    def test_generate_qa_pair_endpoint_returns_data(self):
        request = QARequest(
            request_id="req-001",
            dashboard_id="DB001",
            card_ids=["C0001", "C0002"],
            bbk_id="BBK002",
            user_request="测试问题",
        )

        response_model = asyncio.run(generate_qa_pair(request))

        self.assertEqual(response_model.returnCode, "SUCCESS")
        body = response_model.body
        self.assertEqual(body.dashboard_id, request.dashboard_id)
        self.assertTrue(set(body.card_ids).issubset(set(request.card_ids)))
        self.assertGreaterEqual(len(body.qa_pairs), 1)
        first = body.qa_pairs[0]
        self.assertEqual(first.dashboard_id, request.dashboard_id)
        self.assertIn(first.card_id, request.card_ids)

        log_dir = Path("app/log/test_results")
        log_dir.mkdir(parents=True, exist_ok=True)
        outfile = log_dir / "generate_qa_pair_response.txt"
        outfile.write_text(response_model.model_dump_json(ensure_ascii=False, indent=2))


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