117 lines
3.3 KiB
Python
117 lines
3.3 KiB
Python
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)
|