Files
order-bak/scripts/seed_test_data.py
tech c4290b830e refactor: read config from file, randomize users and dates in seed script
Replace hardcoded connection string with config file loading via
load_config(). Users now randomly chosen from user001-user1000, dates
randomly sampled between 2026-01-01 and UTC now.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-04-22 00:04:32 +08:00

116 lines
4.1 KiB
Python

"""Seed test data into Azure Table emulator for integration testing."""
import random
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
from azure.data.tables import TableServiceClient
from order_bak.config import load_config
TABLE_NAME = "order"
USER_COUNT = 1000
ENTRIES_PER_USER_PER_DAY = 2
PAYED_PER_USER_PER_DAY = 1
DATE_START = datetime(2026, 1, 1, tzinfo=timezone.utc)
DATE_END = datetime.now(timezone.utc)
def seed(config_path: str) -> None:
cfg = load_config(Path(config_path))
service = TableServiceClient.from_connection_string(cfg.azure_connection_string)
service.create_table_if_not_exists(table_name=TABLE_NAME)
client = service.get_table_client(table_name=TABLE_NAME)
# Clear existing data
existing = list(client.list_entities(select=["PartitionKey", "RowKey"]))
by_partition: dict[str, list[dict]] = {}
for e in existing:
pk = e["PartitionKey"]
by_partition.setdefault(pk, []).append(
{"PartitionKey": pk, "RowKey": e["RowKey"]}
)
for pk, entities in by_partition.items():
for i in range(0, len(entities), 100):
batch = entities[i : i + 100]
ops = [("delete", e) for e in batch]
client.submit_transaction(ops)
if existing:
print(f"Cleared {len(existing)} existing records")
# Generate random users and dates
users = [f"user{i:03d}" for i in random.sample(range(1, USER_COUNT + 1), k=min(USER_COUNT, USER_COUNT))]
date_range_seconds = int((DATE_END - DATE_START).total_seconds())
# Seed new data
entities = []
order_idx = 0
for _ in range(len(users)):
user = random.choice(users)
# Random timestamp in [DATE_START, DATE_END)
offset = random.randint(0, max(date_range_seconds - 1, 0))
dt = DATE_START + timedelta(seconds=offset)
# Received orders (state=1)
for _ in range(ENTRIES_PER_USER_PER_DAY):
entities.append({
"PartitionKey": user,
"RowKey": f"test-order-{order_idx:04d}",
"OrderTime": dt.replace(hour=8 + (order_idx % 12), minute=order_idx % 60),
"ReceiveTime": dt.replace(hour=9 + (order_idx % 12), minute=order_idx % 60),
"revenue": 10.0 + order_idx,
"platformOrder": f"plat-{order_idx:04d}",
"prodCount": 1 + (order_idx % 5),
"pfId": user,
"prodId": f"prod-{order_idx % 3}",
"prodPrice": 10.0 + (order_idx % 10),
"state": 1, # Received
"ver": 3,
"customData": None,
})
order_idx += 1
# Payed order (state=0, should be skipped by backup)
entities.append({
"PartitionKey": user,
"RowKey": f"test-order-{order_idx:04d}",
"OrderTime": dt.replace(hour=14, minute=30),
"ReceiveTime": None,
"revenue": 5.0,
"platformOrder": f"plat-payed-{order_idx:04d}",
"prodCount": 1,
"pfId": user,
"prodId": "prod-pay",
"prodPrice": 5.0,
"state": 0, # Payed
"ver": 3,
"customData": None,
})
order_idx += 1
# Insert in batches grouped by PartitionKey
total = 0
by_pk: dict[str, list[dict]] = {}
for e in entities:
by_pk.setdefault(e["PartitionKey"], []).append(e)
for pk, pk_entities in by_pk.items():
for i in range(0, len(pk_entities), 100):
batch = pk_entities[i : i + 100]
ops = [("create", e) for e in batch]
client.submit_transaction(ops)
total += len(batch)
print(f"Seeded {total} orders across {len(by_pk)} partitions")
print(f" Received (state=1): {len([e for e in entities if e['state'] == 1])}")
print(f" Payed (state=0): {len([e for e in entities if e['state'] == 0])}")
if __name__ == "__main__":
config_path = sys.argv[1] if len(sys.argv) > 1 else str(
Path.home() / ".config" / "order-bak" / "config.toml"
)
seed(config_path)