sql_qa_pair/tests/test_generate_operator.py

import asyncio
import unittest

from app.operators.get_dashboard_info import get_dashboard_info
from app.operators.generate_qa_pair import generate_qa_pair


class GenerateOperatorTest(unittest.TestCase):
    def setUp(self):
        self.input_args = {
            "dashboard_id": "DB001",
            "card_ids": ["C0001", "C0002"],
            "bbk": "BBK002",
        }
        self.dashboard_info = get_dashboard_info(self.input_args)

    def test_generate_qa_pair_fallback_output(self):
        result = asyncio.run(
            generate_qa_pair({
                "get_dashboard_info": self.dashboard_info,
                "dashboard_id": self.input_args["dashboard_id"],
                "card_ids": self.input_args["card_ids"],
                "bbk": self.input_args["bbk"],
                "user_request": "",
            })
        )

        self.assertEqual(len(result), len(self.dashboard_info.cards))
        first = result[0]
        print(first)
        self.assertEqual(first.dashboard_id, self.input_args["dashboard_id"])
        self.assertIn(first.card_id, self.input_args["card_ids"])
        self.assertTrue(first.question)


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