from __future__ import annotations from dataclasses import dataclass from decimal import Decimal from app.extensions import db from app.models import Account, Household, HouseholdMember, OptionItem DEMO_ADMIN_USERNAME = "admin" DEFAULT_STRESS_HOUSEHOLD_COUNT = 1000 DEFAULT_STRESS_MEMBERS_PER_HOUSEHOLD = 4 DEFAULT_STRESS_HOUSEHOLD_CODE_PREFIX = "STRESS-" _STRESS_SIDE_OPTIONS: tuple[str, ...] = ("groom_side", "bride_side", "both_side") _STRESS_RELATION_CATEGORY_CODES: tuple[str, ...] = ("relative", "friend", "colleague") _STRESS_TAG_CODE_SETS: tuple[tuple[str, ...], ...] = ((),) _STRESS_INVITE_STATUS_OPTIONS: tuple[str, ...] = ("not_invited", "invited", "confirmed") _STRESS_ATTENDANCE_STATUS_OPTIONS: tuple[str, ...] = ("pending", "attending", "partial", "absent") _STRESS_GIVEN_STATUS_OPTIONS: tuple[str, ...] = ("not_given", "given") _STRESS_CHILD_RED_PACKET_STATUS_OPTIONS: tuple[str, ...] = ("not_given", "partial", "given") _STRESS_GIFT_METHOD_CODES: tuple[str, ...] = ("cash", "transfer") _STRESS_GIFT_SCENE_CODES: tuple[str, ...] = ("pre_wedding", "wedding_day", "post_wedding") _STRESS_RELATION_TO_HEAD_OPTIONS: tuple[str, ...] = ( "本人", "配偶", "子女", "亲属", "同事", "同学", ) _STRESS_GENDER_OPTIONS: tuple[str, ...] = ("male", "female") @dataclass(frozen=True) class DemoAccountSpec: username: str display_name: str password: str role: str @dataclass(frozen=True) class DemoMemberSpec: name: str relation_to_head: str | None gender: str | None age_group: str | None is_child: bool needs_red_packet: bool expected_to_attend: bool actually_attended: bool sort_order: int note: str | None @dataclass(frozen=True) class DemoHouseholdSpec: household_code: str head_name: str phone: str | None side: str relation_category_code: str | None relation_detail_code: str | None tag_codes: tuple[str, ...] invite_status: str attendance_status: str expected_attendee_count: int actual_attendee_count: int child_count: int red_packet_child_count: int total_gift_amount: Decimal gift_method_code: str | None gift_scene_code: str | None favor_status: str candy_status: str child_red_packet_status: str note: str | None members: tuple[DemoMemberSpec, ...] DEMO_ACCOUNT_SPECS: tuple[DemoAccountSpec, ...] = ( DemoAccountSpec( username="editor-demo", display_name="演示录入员", password="EditorDemo123!", role="editor", ), DemoAccountSpec( username="entry-demo", display_name="长辈演示账号", password="EntryDemo123!", role="entry_only", ), DemoAccountSpec( username="quick-editor-demo", display_name="快速录入账号", password="QuickEditor123!", role="quick_editor", ), ) DEMO_HOUSEHOLD_SPECS: tuple[DemoHouseholdSpec, ...] = ( DemoHouseholdSpec( household_code="D001", head_name="王大爷", phone="13800001234", side="groom_side", relation_category_code="relative", relation_detail_code=None, tag_codes=(), invite_status="confirmed", attendance_status="pending", expected_attendee_count=4, actual_attendee_count=0, child_count=1, red_packet_child_count=1, total_gift_amount=Decimal("666.00"), gift_method_code="cash", gift_scene_code="wedding_day", favor_status="not_given", candy_status="given", child_red_packet_status="not_given", note="演示数据：待电话确认当天到场时间。", members=( DemoMemberSpec( name="王大爷", relation_to_head="本人", gender="male", age_group="elder", is_child=False, needs_red_packet=False, expected_to_attend=True, actually_attended=False, sort_order=10, note="户主", ), DemoMemberSpec( name="王阿姨", relation_to_head="配偶", gender="female", age_group="elder", is_child=False, needs_red_packet=False, expected_to_attend=True, actually_attended=False, sort_order=20, note=None, ), DemoMemberSpec( name="王小宝", relation_to_head="孙辈", gender="male", age_group="child", is_child=True, needs_red_packet=True, expected_to_attend=True, actually_attended=False, sort_order=30, note="需准备儿童红包", ), ), ), DemoHouseholdSpec( household_code="D002", head_name="李阿姨", phone="13900004567", side="bride_side", relation_category_code="friend", relation_detail_code=None, tag_codes=(), invite_status="invited", attendance_status="attending", expected_attendee_count=2, actual_attendee_count=2, child_count=0, red_packet_child_count=0, total_gift_amount=Decimal("888.00"), gift_method_code="transfer", gift_scene_code="pre_wedding", favor_status="given", candy_status="given", child_red_packet_status="not_given", note="演示数据：已确认提前到场帮忙。", members=( DemoMemberSpec( name="李阿姨", relation_to_head="本人", gender="female", age_group="adult", is_child=False, needs_red_packet=False, expected_to_attend=True, actually_attended=True, sort_order=10, note=None, ), DemoMemberSpec( name="李叔叔", relation_to_head="配偶", gender="male", age_group="adult", is_child=False, needs_red_packet=False, expected_to_attend=True, actually_attended=True, sort_order=20, note=None, ), ), ), DemoHouseholdSpec( household_code="D003", head_name="陈老师", phone="13700007890", side="both_side", relation_category_code="colleague", relation_detail_code=None, tag_codes=(), invite_status="confirmed", attendance_status="partial", expected_attendee_count=3, actual_attendee_count=2, child_count=1, red_packet_child_count=0, total_gift_amount=Decimal("520.00"), gift_method_code="cash", gift_scene_code="wedding_day", favor_status="given", candy_status="not_given", child_red_packet_status="partial", note="演示数据：一位孩子不单独领红包。", members=( DemoMemberSpec( name="陈老师", relation_to_head="本人", gender="female", age_group="adult", is_child=False, needs_red_packet=False, expected_to_attend=True, actually_attended=True, sort_order=10, note=None, ), DemoMemberSpec( name="陈先生", relation_to_head="配偶", gender="male", age_group="adult", is_child=False, needs_red_packet=False, expected_to_attend=True, actually_attended=True, sort_order=20, note=None, ), DemoMemberSpec( name="陈小朋友", relation_to_head="子女", gender="female", age_group="child", is_child=True, needs_red_packet=False, expected_to_attend=True, actually_attended=False, sort_order=30, note="部分到场示例", ), ), ), ) def seed_demo_accounts() -> int: created_count = 0 for spec in DEMO_ACCOUNT_SPECS: existing_account = db.session.execute( db.select(Account).where(Account.username == spec.username), ).scalar_one_or_none() if existing_account is not None: continue account = Account( username=spec.username, display_name=spec.display_name, role=spec.role, status="active", ) account.set_password(spec.password) db.session.add(account) created_count += 1 return created_count def seed_demo_households() -> int: created_count = 0 option_index = _build_option_index() admin_account = _get_admin_account() admin_id = admin_account.id if admin_account is not None else None for spec in DEMO_HOUSEHOLD_SPECS: existing_household = db.session.execute( db.select(Household).where(Household.household_code == spec.household_code), ).scalar_one_or_none() if existing_household is not None: continue household = Household( household_code=spec.household_code, head_name=spec.head_name, phone=spec.phone, side=spec.side, relation_category_option_id=_resolve_option_id(option_index, "relation_category", spec.relation_category_code), relation_detail_option_id=_resolve_option_id(option_index, "relation_detail", spec.relation_detail_code), tag_option_ids_json=[ option_id for option_id in ( _resolve_option_id(option_index, "tag", tag_code) for tag_code in spec.tag_codes ) if option_id is not None ], invite_status=spec.invite_status, attendance_status=spec.attendance_status, expected_attendee_count=spec.expected_attendee_count, actual_attendee_count=spec.actual_attendee_count, child_count=spec.child_count, red_packet_child_count=spec.red_packet_child_count, total_gift_amount=spec.total_gift_amount, gift_method_option_id=_resolve_option_id(option_index, "gift_method", spec.gift_method_code), gift_scene_option_id=_resolve_option_id(option_index, "gift_scene", spec.gift_scene_code), favor_status=spec.favor_status, candy_status=spec.candy_status, child_red_packet_status=spec.child_red_packet_status, note=spec.note, version=1, created_by=admin_id, updated_by=admin_id, ) db.session.add(household) db.session.flush() for member_spec in spec.members: member = HouseholdMember( household_id=household.id, name=member_spec.name, relation_to_head=member_spec.relation_to_head, gender=member_spec.gender, age_group=member_spec.age_group, is_child=member_spec.is_child, needs_red_packet=member_spec.needs_red_packet, expected_to_attend=member_spec.expected_to_attend, actually_attended=member_spec.actually_attended, sort_order=member_spec.sort_order, note=member_spec.note, created_by=admin_id, updated_by=admin_id, ) db.session.add(member) created_count += 1 return created_count def seed_stress_households( *, household_count: int = DEFAULT_STRESS_HOUSEHOLD_COUNT, members_per_household: int = DEFAULT_STRESS_MEMBERS_PER_HOUSEHOLD, household_code_prefix: str = DEFAULT_STRESS_HOUSEHOLD_CODE_PREFIX, ) -> tuple[int, int]: normalized_prefix = household_code_prefix.strip().upper() if household_count <= 0: raise ValueError("household_count must be greater than 0") if members_per_household <= 0: raise ValueError("members_per_household must be greater than 0") if not normalized_prefix: raise ValueError("household_code_prefix cannot be empty") created_household_count = 0 created_member_count = 0 option_index = _build_option_index() admin_account = _get_admin_account() admin_id = admin_account.id if admin_account is not None else None for index in range(1, household_count + 1): created = _create_stress_household( index=index, household_code=f"{normalized_prefix}{index:04d}", household_code_prefix=normalized_prefix, members_per_household=members_per_household, option_index=option_index, admin_id=admin_id, ) if created is None: continue created_household_count += 1 created_member_count += created return (created_household_count, created_member_count) def _create_stress_household( *, index: int, household_code: str, household_code_prefix: str, members_per_household: int, option_index: dict[tuple[str, str], int], admin_id: int | None, ) -> int | None: existing_household = db.session.execute( db.select(Household).where(Household.household_code == household_code), ).scalar_one_or_none() if existing_household is not None: return None household = _build_stress_household( index=index, household_code=household_code, household_code_prefix=household_code_prefix, members_per_household=members_per_household, option_index=option_index, admin_id=admin_id, ) db.session.add(household) db.session.flush() for member_index in range(1, members_per_household + 1): member = _build_stress_member( household_id=household.id, household_index=index, member_index=member_index, members_per_household=members_per_household, admin_id=admin_id, ) db.session.add(member) return members_per_household def _build_stress_household( *, index: int, household_code: str, household_code_prefix: str, members_per_household: int, option_index: dict[tuple[str, str], int], admin_id: int | None, ) -> Household: relation_category_code = _STRESS_RELATION_CATEGORY_CODES[(index - 1) % len(_STRESS_RELATION_CATEGORY_CODES)] side = _STRESS_SIDE_OPTIONS[(index - 1) % len(_STRESS_SIDE_OPTIONS)] invite_status = _STRESS_INVITE_STATUS_OPTIONS[(index - 1) % len(_STRESS_INVITE_STATUS_OPTIONS)] attendance_status = _STRESS_ATTENDANCE_STATUS_OPTIONS[(index - 1) % len(_STRESS_ATTENDANCE_STATUS_OPTIONS)] tag_codes = _STRESS_TAG_CODE_SETS[(index - 1) % len(_STRESS_TAG_CODE_SETS)] child_count = 1 if members_per_household >= 3 and index % 3 == 0 else 0 red_packet_child_count = child_count if index % 2 == 0 else 0 actual_attendee_count = _stress_actual_attendee_count( expected_attendee_count=members_per_household, attendance_status=attendance_status, ) total_gift_amount = Decimal(index * 50).quantize(Decimal("0.01")) phone_number = f"139{index:08d}"[-11:] normalized_phone = f"1{phone_number[-10:]}" head_name_prefix = "" if household_code_prefix == DEFAULT_STRESS_HOUSEHOLD_CODE_PREFIX else household_code_prefix return Household( household_code=household_code, head_name=f"{head_name_prefix}压力户主{index:04d}", phone=normalized_phone, side=side, relation_category_option_id=_resolve_option_id(option_index, "relation_category", relation_category_code), relation_detail_option_id=None, tag_option_ids_json=[ option_id for option_id in ( _resolve_option_id(option_index, "tag", tag_code) for tag_code in tag_codes ) if option_id is not None ], invite_status=invite_status, attendance_status=attendance_status, expected_attendee_count=members_per_household, actual_attendee_count=actual_attendee_count, child_count=child_count, red_packet_child_count=red_packet_child_count, total_gift_amount=total_gift_amount, gift_method_option_id=_resolve_option_id( option_index, "gift_method", _STRESS_GIFT_METHOD_CODES[(index - 1) % len(_STRESS_GIFT_METHOD_CODES)], ), gift_scene_option_id=_resolve_option_id( option_index, "gift_scene", _STRESS_GIFT_SCENE_CODES[(index - 1) % len(_STRESS_GIFT_SCENE_CODES)], ), favor_status=_STRESS_GIVEN_STATUS_OPTIONS[index % len(_STRESS_GIVEN_STATUS_OPTIONS)], candy_status=_STRESS_GIVEN_STATUS_OPTIONS[(index + 1) % len(_STRESS_GIVEN_STATUS_OPTIONS)], child_red_packet_status=_STRESS_CHILD_RED_PACKET_STATUS_OPTIONS[ (index - 1) % len(_STRESS_CHILD_RED_PACKET_STATUS_OPTIONS) ], note=f"压力测试数据 #{index:04d}", version=1, created_by=admin_id, updated_by=admin_id, ) def _build_stress_member( *, household_id: int, household_index: int, member_index: int, members_per_household: int, admin_id: int | None, ) -> HouseholdMember: is_head = member_index == 1 is_child = members_per_household >= 3 and member_index == members_per_household and household_index % 3 == 0 attended = household_index % 4 in {1, 2} and not is_child return HouseholdMember( household_id=household_id, name=f"压力成员{household_index:04d}-{member_index}", relation_to_head=_stress_relation_to_head(member_index=member_index), gender=_STRESS_GENDER_OPTIONS[(household_index + member_index - 2) % len(_STRESS_GENDER_OPTIONS)], age_group="child" if is_child else _stress_age_group(member_index=member_index, is_head=is_head), is_child=is_child, needs_red_packet=is_child and household_index % 2 == 0, expected_to_attend=True, actually_attended=attended, sort_order=member_index * 10, note="压力测试儿童成员" if is_child else ("户主" if is_head else None), created_by=admin_id, updated_by=admin_id, ) def _stress_actual_attendee_count(*, expected_attendee_count: int, attendance_status: str) -> int: if attendance_status == "attending": return expected_attendee_count if attendance_status == "partial": return max(expected_attendee_count - 1, 1) if attendance_status == "absent": return 0 return 0 def _stress_relation_to_head(*, member_index: int) -> str: if member_index == 1: return "本人" if member_index == 2: return "配偶" if member_index == 3: return "子女" return _STRESS_RELATION_TO_HEAD_OPTIONS[(member_index - 1) % len(_STRESS_RELATION_TO_HEAD_OPTIONS)] def _stress_age_group(*, member_index: int, is_head: bool) -> str: if is_head: return "adult" if member_index == 2: return "adult" return "elder" if member_index % 2 == 0 else "adult" def demo_login_lines() -> tuple[str, ...]: return ( "admin / ChangeMe123!（管理员）", "editor-demo / EditorDemo123!（演示录入员）", "entry-demo / EntryDemo123!（长辈录入演示）", "quick-editor-demo / QuickEditor123!（快速录入账号）", ) def _build_option_index() -> dict[tuple[str, str], int]: option_items = db.session.execute(db.select(OptionItem)).scalars().all() return { (option_item.option_group, option_item.option_code): option_item.id for option_item in option_items } def _resolve_option_id( option_index: dict[tuple[str, str], int], option_group: str, option_code: str | None, ) -> int | None: if option_code is None: return None return option_index.get((option_group, option_code)) def _get_admin_account() -> Account | None: return db.session.execute( db.select(Account).where(Account.username == DEMO_ADMIN_USERNAME), ).scalar_one_or_none()