zjpb.net/test_site_creation.py

"""测试网站创建功能 - 验证code和slug自动生成"""
import sys
import os

# 设置UTF-8编码
if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

from app import create_app
from models import db, Site

def test_site_creation_with_autoflush():
    """测试通过Flask-Admin添加网站时的autoflush问题"""
    print("=" * 60)
    print("测试：验证on_model_change中的no_autoflush修复")
    print("=" * 60)

    app = create_app()

    with app.app_context():
        try:
            # 模拟Flask-Admin创建网站的过程
            print("\n1. 创建新网站对象（code和slug为空）...")
            site = Site(
                name='测试网站WaytoAGI',
                url='https://test.waytoagi.com',
                short_desc='这是一个测试网站',
                description='用于测试code和slug自动生成功能',
                is_active=True
            )

            print(f"   初始状态: code={site.code}, slug={site.slug}")

            # 添加到session
            print("\n2. 添加到数据库session...")
            db.session.add(site)

            # 模拟on_model_change的调用
            print("\n3. 执行on_model_change逻辑（生成code和slug）...")
            import re
            import random
            from pypinyin import lazy_pinyin

            # 使用no_autoflush（这是修复的关键）
            with db.session.no_autoflush:
                # 生成code
                if not site.code:
                    for _ in range(10):
                        code = str(random.randint(10000000, 99999999))
                        existing = Site.query.filter(Site.code == code).first()
                        if not existing:
                            site.code = code
                            print(f"   生成code: {code}")
                            break

                # 生成slug
                if not site.slug:
                    slug = ''.join(lazy_pinyin(site.name))
                    slug = slug.lower()
                    slug = re.sub(r'[^\w\s-]', '', slug)
                    slug = re.sub(r'[-\s]+', '-', slug).strip('-')

                    if not slug:
                        slug = f"site-{site.code}"

                    base_slug = slug[:50]
                    counter = 1
                    final_slug = slug

                    while counter < 100:
                        existing = Site.query.filter(Site.slug == final_slug).first()
                        if not existing:
                            break
                        final_slug = f"{base_slug}-{counter}"
                        counter += 1

                    site.slug = final_slug
                    print(f"   生成slug: {final_slug}")

            print(f"\n4. 提交到数据库...")
            print(f"   最终状态: code={site.code}, slug={site.slug}")

            # 提交事务
            db.session.commit()

            print("\n[SUCCESS] 网站创建成功！")
            print(f"   ID: {site.id}")
            print(f"   名称: {site.name}")
            print(f"   Code: {site.code}")
            print(f"   Slug: {site.slug}")

            # 验证数据
            print("\n5. 验证数据完整性...")
            assert site.code is not None, "code不能为空"
            assert site.slug is not None, "slug不能为空"
            assert len(site.code) == 8, "code必须是8位数字"
            assert site.code.isdigit(), "code必须是纯数字"
            print("   [OK] 所有验证通过")

            # 清理测试数据
            print("\n6. 清理测试数据...")
            db.session.delete(site)
            db.session.commit()
            print("   [OK] 测试数据已清理")

            return True

        except Exception as e:
            db.session.rollback()
            print(f"\n[FAILED] 测试失败：{str(e)}")
            import traceback
            traceback.print_exc()
            return False

def test_unique_code_generation():
    """测试code唯一性生成"""
    print("\n" + "=" * 60)
    print("测试：验证code唯一性检查")
    print("=" * 60)

    app = create_app()

    with app.app_context():
        try:
            import random

            print("\n1. 生成测试code...")
            test_code = str(random.randint(10000000, 99999999))
            print(f"   测试code: {test_code}")

            print("\n2. 创建第一个网站...")
            site1 = Site(
                code=test_code,
                name='测试网站1',
                url='https://test1.com',
                slug='test-site-1',
                is_active=True
            )
            db.session.add(site1)
            db.session.commit()
            print(f"   [OK] 网站1创建成功，code={site1.code}")

            print("\n3. 尝试在no_autoflush中查询该code...")
            with db.session.no_autoflush:
                existing = Site.query.filter(Site.code == test_code).first()
                if existing:
                    print(f"   [OK] 成功查询到已存在的code: {existing.code}")
                else:
                    print(f"   [ERROR] 未能查询到已存在的code")
                    return False

            print("\n4. 清理测试数据...")
            db.session.delete(site1)
            db.session.commit()
            print("   [OK] 测试数据已清理")

            return True

        except Exception as e:
            db.session.rollback()
            print(f"\n[FAILED] 测试失败：{str(e)}")
            import traceback
            traceback.print_exc()
            return False

if __name__ == "__main__":
    print("\n开始测试网站创建功能...\n")

    test1_passed = test_site_creation_with_autoflush()
    test2_passed = test_unique_code_generation()

    print("\n" + "=" * 60)
    print("测试结果汇总")
    print("=" * 60)
    print(f"测试1 - autoflush修复验证: {'[PASSED]' if test1_passed else '[FAILED]'}")
    print(f"测试2 - code唯一性验证: {'[PASSED]' if test2_passed else '[FAILED]'}")

    if test1_passed and test2_passed:
        print("\n[SUCCESS] 所有测试通过！on_model_change修复有效。")
        sys.exit(0)
    else:
        print("\n[FAILED] 部分测试失败，需要进一步检查。")
        sys.exit(1)