from unittest.mock import patch

from django.db.models.query import QuerySet
from django.test import SimpleTestCase, TestCase

from apps.develop.models import Project
from apps.mtm.models import Material
from apps.system.models import User
from apps.wpmw.models import WprDefect


class SnowflakeBulkCreateTests(SimpleTestCase):
    @patch.object(QuerySet, "bulk_create")
    @patch(
        "apps.utils.models.idWorker.get_id",
        side_effect=[100, 101],
    )
    def test_base_model_bulk_create_assigns_missing_ids(
        self,
        get_id,
        bulk_create,
    ):
        objs = [
            WprDefect(defect_id="10", wpr_id="20"),
            WprDefect(defect_id="11", wpr_id="20"),
        ]

        WprDefect.objects.bulk_create(objs)

        self.assertEqual([obj.id for obj in objs], ["100", "101"])
        self.assertEqual(get_id.call_count, 2)
        bulk_create.assert_called_once()

    @patch.object(QuerySet, "bulk_create")
    @patch("apps.utils.models.idWorker.get_id")
    def test_base_model_bulk_create_preserves_existing_ids(
        self,
        get_id,
        bulk_create,
    ):
        objs = [
            WprDefect(id=100, defect_id="10", wpr_id="20"),
            WprDefect(id="101", defect_id="11", wpr_id="20"),
        ]

        WprDefect.objects.bulk_create(objs)

        self.assertEqual([obj.id for obj in objs], ["100", "101"])
        get_id.assert_not_called()
        bulk_create.assert_called_once()

    @patch.object(QuerySet, "bulk_create")
    @patch(
        "apps.utils.models.idWorker.get_id",
        side_effect=[200, 201],
    )
    def test_soft_delete_manager_uses_snowflake_bulk_create(
        self,
        get_id,
        bulk_create,
    ):
        objs = [
            Material(name="物料1"),
            Material(name="物料2"),
        ]

        Material.objects.bulk_create(objs)

        self.assertEqual([obj.id for obj in objs], ["200", "201"])
        self.assertEqual(get_id.call_count, 2)
        bulk_create.assert_called_once()

    @patch.object(QuerySet, "bulk_create")
    @patch(
        "apps.utils.models.idWorker.get_id",
        side_effect=[300, 301],
    )
    def test_custom_user_manager_uses_snowflake_bulk_create(
        self,
        get_id,
        bulk_create,
    ):
        objs = [
            User(username="bulk-user-1"),
            User(username="bulk-user-2"),
        ]

        User.objects.bulk_create(objs)

        self.assertEqual([obj.id for obj in objs], ["300", "301"])
        self.assertEqual(get_id.call_count, 2)
        bulk_create.assert_called_once()


class SnowflakeBulkCreateDatabaseTests(TestCase):
    @patch(
        "apps.utils.models.idWorker.get_id",
        side_effect=[900000000000000001, 900000000000000002],
    )
    def test_bulk_create_persists_generated_ids(self, get_id):
        objs = [
            Project(name="批量项目1", code="bulk-project-1"),
            Project(name="批量项目2", code="bulk-project-2"),
        ]

        Project.objects.bulk_create(objs)

        self.assertEqual(
            [obj.id for obj in objs],
            ["900000000000000001", "900000000000000002"],
        )
        self.assertEqual(
            set(Project.objects.values_list("id", flat=True)),
            {"900000000000000001", "900000000000000002"},
        )
        self.assertEqual(get_id.call_count, 2)