tests/test_query_data.py

import unittest

from core.query_data import (
    normalize_query_text,
    sanitize_description_queries,
    sanitize_query_rows,
)


class QueryDataTests(unittest.TestCase):
    def test_normalization(self):
        self.assertEqual(normalize_query_text("  잔고   조회? "), "잔고 조회")

    def test_conflicting_queries_are_dropped(self):
        rows = [
            {"query": "잔고 조회", "menu_id": "SCR_1"},
            {"query": "잔고 조회?", "menu_id": "SCR_2"},
            {"query": "미체결 확인", "menu_id": "SCR_3"},
        ]
        clean, stats = sanitize_query_rows(rows)
        self.assertEqual(
            clean,
            [{"query": "미체결 확인", "menu_id": "SCR_3"}],
        )
        self.assertEqual(stats["conflicting"], 2)

    def test_held_out_and_duplicates_are_dropped(self):
        rows = [
            {"query": "OTP 등록", "menu_id": "SCR_1"},
            {"query": "OTP 등록", "menu_id": "SCR_1"},
            {"query": "환전하기", "menu_id": "SCR_2"},
        ]
        clean, stats = sanitize_query_rows(
            rows,
            excluded_queries={"환전하기"},
        )
        self.assertEqual(len(clean), 1)
        self.assertEqual(stats["duplicate"], 1)
        self.assertEqual(stats["held_out"], 1)

    def test_explicitly_shared_queries_are_retained_for_multiple_menus(self):
        rows = [
            {
                "query": "내 계좌 보여줘",
                "menu_id": "SCR_1300",
                "allow_shared_query": True,
            },
            {
                "query": "내 계좌 보여줘",
                "menu_id": "SCR_3075",
                "allow_shared_query": True,
            },
        ]

        clean, stats = sanitize_query_rows(rows)

        self.assertEqual(len(clean), 2)
        self.assertEqual(stats["conflicting"], 0)
        self.assertEqual(stats["shared_conflicting_texts"], 1)

    def test_description_source_is_not_mutated(self):
        source = [
            {"menu_id": "SCR_1", "sample_queries": ["잔고 조회", "OTP 등록"]},
            {"menu_id": "SCR_2", "sample_queries": ["잔고 조회"]},
        ]
        clean, stats = sanitize_description_queries(source)
        self.assertEqual(source[0]["sample_queries"], ["잔고 조회", "OTP 등록"])
        self.assertEqual(clean[0]["sample_queries"], ["OTP 등록"])
        self.assertEqual(clean[1]["sample_queries"], [])
        self.assertEqual(stats["removed_conflicting"], 2)


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