from django.db import connection
from django.test import TestCase
from django.test.utils import CaptureQueriesContext

from apps.mtm.models import Material
from apps.qm.models import Defect
from apps.system.models import Dept, User
from apps.wpm.models import Handover, WMaterial
from apps.wpm.serializers import HandoverSerializer


class HandoverSerializerQueryTests(TestCase):
    @classmethod
    def setUpTestData(cls):
        cls.dept = Dept.objects.create(name="合批查询测试车间")
        cls.user = User.objects.create_user(username="handover-query-user")
        cls.material = Material.objects.create(name="合批查询测试物料")
        cls.defect_b = Defect.objects.create(
            name="合批查询测试B类缺陷",
            cate=Defect.cate_list[0],
            okcate=Defect.DEFECT_OK_B,
        )
        cls.inventories = [
            WMaterial.objects.create(
                material=cls.material,
                batch=f"QUERY-{index}",
                count=1,
                state=WMaterial.WM_OK,
                defect=cls.defect_b if index == 2 else None,
                belong_dept=cls.dept,
            )
            for index in range(3)
        ]

    def test_handover_inventory_is_loaded_in_one_query(self):
        serializer = HandoverSerializer(data={
            "send_date": "2026-08-07",
            "send_user": self.user.id,
            "send_dept": self.dept.id,
            "recive_dept": self.dept.id,
            "handoverb": [
                {"wm": inventory.id, "count": 1}
                for inventory in self.inventories
            ],
            "new_batch": "QUERY-MERGED",
            "clear_defect": True,
            "type": Handover.H_NORMAL,
            "mtype": Handover.H_MERGE,
        })

        with CaptureQueriesContext(connection) as queries:
            self.assertTrue(serializer.is_valid(), serializer.errors)

        inventory_queries = [
            query["sql"] for query in queries
            if 'FROM "wpm_wmaterial"' in query["sql"]
        ]
        self.assertEqual(len(inventory_queries), 1, inventory_queries)