#!/usr/bin/env python3
import re

# 1. Fix internal/service/account/service_test.go
print("Fixing service_test.go...")
with open('internal/service/account/service_test.go', 'r') as f:
    content = f.read()

content = re.sub(
    r'(\taccountRoleStore := postgres\.NewAccountRoleStore\(tx, rdb\)\n)',
    r'\1\tshopRoleStore := postgres.NewShopRoleStore(tx, rdb)\n',
    content,
    count=1
)

content = re.sub(
    r'New\(accountStore, roleStore, accountRoleStore, (&MockShopStore{[^}]*}), (&MockEnterpriseStore{[^}]*}), (&MockAuditService{[^}]*})\)',
    r'New(accountStore, roleStore, accountRoleStore, nil, \1, \2, \3)',
    content
)

with open('internal/service/account/service_test.go', 'w') as f:
    f.write(content)
print("✓ Fixed service_test.go")

# 2. Fix tests/unit/permission_check_test.go
print("Fixing permission_check_test.go...")
with open('tests/unit/permission_check_test.go', 'r') as f:
    content = f.read()

content = re.sub(
    r'permission\.New\(permissionStore, accountRoleStore, rolePermStore, rdb\)',
    'permission.New(permissionStore, accountRoleStore, rolePermStore, nil, rdb)',
    content
)

with open('tests/unit/permission_check_test.go', 'w') as f:
    f.write(content)
print("✓ Fixed permission_check_test.go")

# 3. Fix tests/unit/permission_cache_test.go
print("Fixing permission_cache_test.go...")
with open('tests/unit/permission_cache_test.go', 'r') as f:
    content = f.read()

content = re.sub(
    r'permission\.New\(permissionStore, accountRoleStore, rolePermStore, rdb\)',
    'permission.New(permissionStore, accountRoleStore, rolePermStore, nil, rdb)',
    content
)

with open('tests/unit/permission_cache_test.go', 'w') as f:
    f.write(content)
print("✓ Fixed permission_cache_test.go")

# 4. Fix tests/unit/shop_service_test.go
print("Fixing shop_service_test.go...")
with open('tests/unit/shop_service_test.go', 'r') as f:
    content = f.read()

content = re.sub(
    r'(\taccountStore := postgres\.NewAccountStore\(tx, rdb\)\n)',
    r'\1\tshopRoleStore := postgres.NewShopRoleStore(tx, rdb)\n\troleStore := postgres.NewRoleStore(tx)\n',
    content,
    count=1
)

content = re.sub(
    r'shop\.New\(shopStore, accountStore\)',
    'shop.New(shopStore, accountStore, shopRoleStore, roleStore)',
    content
)

with open('tests/unit/shop_service_test.go', 'w') as f:
    f.write(content)
print("✓ Fixed shop_service_test.go")

print("\nAll files fixed!")