deep-risk/backend/init_test_data.py
2025-12-14 20:08:27 +08:00

406 lines
16 KiB
Python

#!/usr/bin/env python3
"""
初始化测试数据
为数据库填充一些示例数据以便开发和测试
"""
import asyncio
import sys
from pathlib import Path
from datetime import date, datetime
from decimal import Decimal
from typing import Optional, List
import uuid
import random
# 添加项目根目录到Python路径
sys.path.insert(0, str(Path(__file__).parent))
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import text
from app.database import AsyncSessionLocal, engine
from app.models.streamer import McnAgency, StreamerInfo, PlatformRecharge
from app.models.contract import Contract
from app.models.order import Order
from app.models.settlement import Settlement
from app.models.expense import Expense
from app.models.tax_declaration import TaxDeclaration
from app.models.bank_transaction import BankTransaction
from app.models.invoice import Invoice
from app.models.risk_detection import DetectionRule
from loguru import logger
async def create_mcn_agencies(session: AsyncSession) -> List[McnAgency]:
"""创建MCN机构数据"""
logger.info("正在创建MCN机构数据...")
agencies_data = [
{
"mcn_id": f"MCN{random.randint(1000, 9999)}",
"mcn_name": "无忧传媒",
"unified_social_credit_code": "91510000MA6CW1234X",
"legal_person_name": "张伟",
"legal_person_id_card": "110101199001011234",
"registered_address": "北京市朝阳区望京SOHO T3座",
"bank_account_no": "6217000012345678901",
"bank_name": "中国建设银行北京望京支行",
"mcn_type": "ecommerce_based",
"status": "active",
"contact_person": "李娜",
"contact_phone": "13800138001",
"contact_email": "contact@wuyoumc.com",
},
{
"mcn_id": f"MCN{random.randint(1000, 9999)}",
"mcn_name": "微播易",
"unified_social_credit_code": "91510000MA6CW5678Y",
"legal_person_name": "王芳",
"legal_person_id_card": "110101199002022345",
"registered_address": "上海市浦东新区陆家嘴金融中心",
"bank_account_no": "6227000012345678902",
"bank_name": "中国工商银行上海陆家嘴支行",
"mcn_type": "tip_based",
"status": "active",
"contact_person": "赵明",
"contact_phone": "13900139002",
"contact_email": "contact@weiboyi.com",
},
{
"mcn_id": f"MCN{random.randint(1000, 9999)}",
"mcn_name": "新榜",
"unified_social_credit_code": "91510000MA6CW9012Z",
"legal_person_name": "刘强",
"legal_person_id_card": "110101199003033456",
"registered_address": "广州市天河区珠江新城IFC",
"bank_account_no": "6228480012345678903",
"bank_name": "中国农业银行广州天河支行",
"mcn_type": "mixed",
"status": "active",
"contact_person": "陈静",
"contact_phone": "13700137003",
"contact_email": "contact@newrank.com",
},
]
agencies = []
for data in agencies_data:
agency = McnAgency(**data)
session.add(agency)
agencies.append(agency)
await session.flush()
logger.info(f"✅ 创建了 {len(agencies)} 个MCN机构")
return agencies
async def create_streamers(session: AsyncSession, mcn_agencies: List[McnAgency]) -> List[StreamerInfo]:
"""创建主播数据"""
logger.info("正在创建主播数据...")
streamers_data = [
{
"streamer_id": f"STR{random.randint(10000, 99999)}",
"streamer_name": "李佳琦",
"entity_type": "enterprise",
"id_card_no": "320101199005054321",
"phone_number": "15800158001",
"unified_social_credit_code": "91320000MA6CW3456A",
"bank_account_no": "6222600260012345678",
"bank_name": "交通银行南京新街口支行",
"wechat_pay_account": "lijiaqi_official",
"alipay_account": "lijiaqi@alipay.com",
"douyin_account": "lijiaqi",
"status": "active",
"registration_date": date(2020, 1, 15),
"mcn_agency_id": mcn_agencies[0].id if mcn_agencies else None,
},
{
"streamer_id": f"STR{random.randint(10000, 99999)}",
"streamer_name": "薇娅",
"entity_type": "enterprise",
"id_card_no": "330101199006065432",
"phone_number": "15900159002",
"unified_social_credit_code": "91330000MA6CW7890B",
"bank_account_no": "6222600260098765432",
"bank_name": "交通银行杭州湖墅支行",
"wechat_pay_account": "viya_official",
"alipay_account": "viya@alipay.com",
"douyin_account": "viya",
"status": "active",
"registration_date": date(2019, 6, 20),
"mcn_agency_id": mcn_agencies[1].id if len(mcn_agencies) > 1 else None,
},
{
"streamer_id": f"STR{random.randint(10000, 99999)}",
"streamer_name": "罗永浩",
"entity_type": "individual",
"id_card_no": "110101197205071234",
"phone_number": "15000150003",
"bank_account_no": "6228480012345678901",
"bank_name": "中国农业银行北京建国门支行",
"wechat_pay_account": "luoyonghao",
"douyin_account": "luoyonghao",
"status": "active",
"registration_date": date(2021, 3, 10),
"mcn_agency_id": mcn_agencies[0].id if mcn_agencies else None,
},
{
"streamer_id": f"STR{random.randint(10000, 99999)}",
"streamer_name": "辛巴",
"entity_type": "sole_proprietor",
"id_card_no": "230101199007082345",
"phone_number": "15700157004",
"bank_account_no": "6217996010012345678",
"bank_name": "中国邮政储蓄银行哈尔滨中央大街支行",
"wechat_pay_account": "xinba_live",
"alipay_account": "xinba@alipay.com",
"douyin_account": "xinba",
"status": "active",
"registration_date": date(2020, 9, 5),
"mcn_agency_id": mcn_agencies[2].id if len(mcn_agencies) > 2 else None,
},
{
"streamer_id": f"STR{random.randint(10000, 99999)}",
"streamer_name": "李子柒",
"entity_type": "individual",
"id_card_no": "510101199008093456",
"phone_number": "15600156005",
"bank_account_no": "6217000012345678901",
"bank_name": "中国建设银行成都蜀都支行",
"wechat_pay_account": "liziqi",
"douyin_account": "liziqi",
"status": "active",
"registration_date": date(2020, 12, 1),
"mcn_agency_id": None,
},
]
streamers = []
for data in streamers_data:
streamer = StreamerInfo(**data)
session.add(streamer)
streamers.append(streamer)
await session.flush()
logger.info(f"✅ 创建了 {len(streamers)} 个主播")
return streamers
async def create_platform_recharges(session: AsyncSession, streamers: List[StreamerInfo]) -> List[PlatformRecharge]:
"""创建平台充值数据"""
logger.info("正在创建平台充值数据...")
recharge_data = []
payment_methods = ["wechat", "alipay", "unionpay", "applepay"]
statuses = ["success", "pending", "failed", "refunded"]
for i in range(50):
streamer = random.choice(streamers)
recharge_data.append({
"recharge_id": f"RCH{uuid.uuid4().hex[:8].upper()}",
"user_id": f"USER{random.randint(10000, 99999)}",
"user_name": streamer.streamer_name,
"user_phone": streamer.phone_number,
"recharge_amount": Decimal(str(random.uniform(100, 10000))),
"recharge_time": datetime.now(),
"payment_method": random.choice(payment_methods),
"transaction_no": f"TXN{uuid.uuid4().hex[:16].upper()}",
"platform_order_no": f"ORD{uuid.uuid4().hex[:12].upper()}",
"actual_amount_cny": Decimal(str(random.uniform(100, 10000))),
"total_coins": Decimal(str(random.uniform(1000, 100000))),
"status": random.choice(statuses),
})
recharges = []
for data in recharge_data:
recharge = PlatformRecharge(**data)
session.add(recharge)
recharges.append(recharge)
await session.flush()
logger.info(f"✅ 创建了 {len(recharges)} 条充值记录")
return recharges
async def create_contracts(session: AsyncSession, streamers: List[StreamerInfo]) -> List[Contract]:
"""创建分成协议数据"""
logger.info("正在创建分成协议数据...")
contracts_data = []
contract_types = ["revenue_share", "service_fee", "guaranteed_minimum"]
revenue_types = ["live_streaming", "advertising", "product_sales"]
for streamer in streamers[:3]:
for i in range(2):
contracts_data.append({
"contract_id": f"CONTRACT{uuid.uuid4().hex[:8].upper()}",
"contract_no": f"HT{uuid.uuid4().hex[:10].upper()}",
"contract_type": random.choice(contract_types),
"streamer_id": streamer.streamer_id,
"revenue_type": random.choice(revenue_types),
"platform_ratio": Decimal(str(random.uniform(10, 50))),
"streamer_ratio": Decimal(str(random.uniform(50, 90))),
"settlement_cycle": random.choice(["monthly", "quarterly", "annually"]),
"contract_start_date": date(2023, 1, 1),
"contract_end_date": date(2025, 12, 31),
"contract_status": "active",
"remark": f"{streamer.streamer_name}的直播合作协议",
})
contracts = []
for data in contracts_data:
contract = Contract(**data)
session.add(contract)
contracts.append(contract)
await session.flush()
logger.info(f"✅ 创建了 {len(contracts)} 份分成协议")
return contracts
async def create_orders(session: AsyncSession, streamers: List[StreamerInfo]) -> List[Order]:
"""创建订单数据"""
logger.info("正在创建订单数据...")
orders_data = []
platforms = ["douyin", "kuaishou", "taobao", "tmall", "xiaohongshu"]
statuses = ["pending", "paid", "shipped", "delivered", "completed", "cancelled"]
for streamer in streamers:
for i in range(10):
orders_data.append({
"order_id": f"ORD{uuid.uuid4().hex[:12].upper()}",
"platform_order_no": f"PL{uuid.uuid4().hex[:16].upper()}",
"ecommerce_platform": random.choice(platforms),
"streamer_id": streamer.streamer_id,
"product_id": f"PROD{random.randint(1000, 9999)}",
"product_name": f"商品名称_{random.randint(1000, 9999)}",
"quantity": random.randint(1, 5),
"original_price": Decimal(str(random.uniform(50, 500))),
"sale_price": Decimal(str(random.uniform(40, 400))),
"total_amount": Decimal(str(random.uniform(100, 2000))),
"actual_payment": Decimal(str(random.uniform(100, 2000))),
"commission_ratio": Decimal(str(random.uniform(5, 30))),
"commission_amount": Decimal(str(random.uniform(10, 200))),
"streamer_commission": Decimal(str(random.uniform(5, 180))),
"buyer_id": f"BUYER{random.randint(1000, 9999)}",
"order_time": datetime.now(),
"order_status": random.choice(statuses),
"is_commission_settled": random.choice([True, False]),
})
orders = []
for data in orders_data:
order = Order(**data)
session.add(order)
orders.append(order)
await session.flush()
logger.info(f"✅ 创建了 {len(orders)} 个订单")
return orders
async def create_detection_rules(session: AsyncSession) -> List[DetectionRule]:
"""创建风险检测规则"""
logger.info("正在创建风险检测规则...")
rules_data = [
{
"rule_id": f"RULE{uuid.uuid4().hex[:8].upper()}",
"rule_name": "大额交易检测",
"algorithm_code": "LARGE_TRANSACTION_CHECK",
"description": "检测单笔交易金额超过阈值的交易",
"parameters": {"threshold": 100000, "time_window": 24},
"is_enabled": True,
},
{
"rule_id": f"RULE{uuid.uuid4().hex[:8].upper()}",
"rule_name": "频繁充值检测",
"algorithm_code": "FREQUENT_RECHARGE_CHECK",
"description": "检测短时间内多次充值的异常行为",
"parameters": {"count_threshold": 10, "time_window": 1},
"is_enabled": True,
},
{
"rule_id": f"RULE{uuid.uuid4().hex[:8].upper()}",
"rule_name": "可疑关联检测",
"algorithm_code": "SUSPICIOUS_RELATIONSHIP_CHECK",
"description": "检测账户之间的可疑关联关系",
"parameters": {"relationship_depth": 2, "transaction_threshold": 50000},
"is_enabled": True,
},
{
"rule_id": f"RULE{uuid.uuid4().hex[:8].upper()}",
"rule_name": "税务异常检测",
"algorithm_code": "TAX_ANOMALY_CHECK",
"description": "检测税务申报数据的异常情况",
"parameters": {"tax_rate_threshold": 0.25, "revenue_variance": 0.3},
"is_enabled": True,
},
]
rules = []
for data in rules_data:
rule = DetectionRule(**data)
session.add(rule)
rules.append(rule)
await session.flush()
logger.info(f"✅ 创建了 {len(rules)} 个风险检测规则")
return rules
async def init_test_data():
"""初始化测试数据"""
logger.info("=" * 60)
logger.info("开始初始化测试数据")
logger.info("=" * 60)
async with AsyncSessionLocal() as session:
try:
# 检查是否已有数据
result = await session.execute(text("SELECT COUNT(*) FROM streamer_info"))
count = result.scalar()
if count > 0:
logger.warning(f"数据库中已有 {count} 条主播数据")
logger.warning("是否继续初始化?(建议先清空数据库)")
# 为了自动化,这里直接继续
logger.warning("跳过检查,继续初始化...")
# 创建数据
mcn_agencies = await create_mcn_agencies(session)
streamers = await create_streamers(session, mcn_agencies)
await create_platform_recharges(session, streamers)
await create_contracts(session, streamers)
await create_orders(session, streamers)
await create_detection_rules(session)
# 提交事务
await session.commit()
logger.info("=" * 60)
logger.info("✅ 测试数据初始化完成!")
logger.info("=" * 60)
except Exception as e:
logger.error(f"初始化测试数据失败: {e}")
await session.rollback()
import traceback
traceback.print_exc()
raise
async def main():
"""主函数"""
try:
await init_test_data()
except Exception as e:
logger.error(f"程序执行失败: {e}")
sys.exit(1)
finally:
await engine.dispose()
if __name__ == "__main__":
asyncio.run(main())