test_auth.py

# test_auth.py
"""
认证系统测试脚本
使用方法：python test_auth.py
"""
import requests
import json
from datetime import datetime

BASE_URL = "http://localhost:8000"

def print_section(title):
    print("\n" + "=" * 60)
    print(f"  {title}")
    print("=" * 60)

def test_register():
    """测试用户注册"""
    print_section("1. 测试用户注册")
    
    test_user = f"test_user_{datetime.now().strftime('%H%M%S')}"
    
    response = requests.post(
        f"{BASE_URL}/register",
        json={"username": test_user, "password": "test123"}
    )
    
    print(f"请求: POST /register")
    print(f"用户名: {test_user}")
    print(f"响应状态: {response.status_code}")
    print(f"响应内容: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")
    
    return test_user, response.json()

def test_login(username, password):
    """测试用户登录"""
    print_section("2. 测试用户登录")
    
    response = requests.post(
        f"{BASE_URL}/login",
        json={"username": username, "password": password}
    )
    
    print(f"请求: POST /login")
    print(f"用户名: {username}")
    print(f"响应状态: {response.status_code}")
    result = response.json()
    print(f"响应内容: {json.dumps(result, ensure_ascii=False, indent=2)}")
    
    if result.get("status") == "success":
        token = result.get("token")
        print(f"\n✅ 登录成功！Token: {token[:50]}...")
        return token
    else:
        print(f"\n❌ 登录失败: {result.get('message')}")
        return None

def test_chat_without_token():
    """测试不带 Token 的对话（应该失败）"""
    print_section("3. 测试无 Token 访问（预期失败）")
    
    try:
        response = requests.post(
            f"{BASE_URL}/chat",
            json={"question": "你好", "username": "admin"}
        )
        print(f"请求: POST /chat (无 Token)")
        print(f"响应状态: {response.status_code}")
        print(f"响应内容: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")
    except Exception as e:
        print(f"请求失败: {e}")

def test_chat_with_token(token, username):
    """测试带 Token 的对话"""
    print_section("4. 测试带 Token 访问（预期成功）")
    
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.post(
        f"{BASE_URL}/chat",
        json={"question": "你好，请介绍一下自己", "username": username},
        headers=headers
    )
    
    print(f"请求: POST /chat (带 Token)")
    print(f"响应状态: {response.status_code}")
    result = response.json()
    
    if response.status_code == 200:
        answer = result.get("answer", "")
        logs = result.get("logs", [])
        print(f"\n✅ 对话成功！")
        print(f"回答: {answer[:100]}...")
        print(f"日志条数: {len(logs)}")
    else:
        print(f"\n❌ 对话失败")
        print(f"响应内容: {json.dumps(result, ensure_ascii=False, indent=2)}")

def test_get_history(token):
    """测试获取历史记录"""
    print_section("5. 测试获取历史记录")
    
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.get(
        f"{BASE_URL}/get_history",
        headers=headers
    )
    
    print(f"请求: GET /get_history")
    print(f"响应状态: {response.status_code}")
    
    if response.status_code == 200:
        history = response.json()
        print(f"\n✅ 获取成功！")
        print(f"历史记录条数: {len(history)}")
        if history:
            print(f"最新记录: {history[-1]}")
    else:
        print(f"响应内容: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")

def test_clear_history(token):
    """测试清空历史记录"""
    print_section("6. 测试清空历史记录")
    
    headers = {"Authorization": f"Bearer {token}"}
    
    response = requests.post(
        f"{BASE_URL}/clear_history",
        headers=headers
    )
    
    print(f"请求: POST /clear_history")
    print(f"响应状态: {response.status_code}")
    result = response.json()
    print(f"响应内容: {json.dumps(result, ensure_ascii=False, indent=2)}")
    
    if result.get("status") == "success":
        print(f"\n✅ 清空成功！")

def test_invalid_token():
    """测试无效 Token"""
    print_section("7. 测试无效 Token（预期失败）")
    
    headers = {"Authorization": "Bearer invalid_token_12345"}
    
    response = requests.post(
        f"{BASE_URL}/chat",
        json={"question": "测试", "username": "test"},
        headers=headers
    )
    
    print(f"请求: POST /chat (无效 Token)")
    print(f"响应状态: {response.status_code}")
    print(f"响应内容: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")

def main():
    print("\n" + "🔒" * 30)
    print("     RAG 系统认证功能测试")
    print("🔒" * 30)
    
    try:
        # 测试 1: 注册新用户
        username, register_result = test_register()
        
        if register_result.get("status") != "success":
            print("\n⚠️  注册失败，使用默认用户 'admin' 进行测试")
            username = "admin"
            password = "123456"
        else:
            password = "test123"
        
        # 测试 2: 登录
        token = test_login(username, password)
        
        if not token:
            print("\n❌ 无法获取 Token，测试中止")
            return
        
        # 测试 3: 不带 Token 访问
        test_chat_without_token()
        
        # 测试 4: 带 Token 对话
        test_chat_with_token(token, username)
        
        # 测试 5: 获取历史记录
        test_get_history(token)
        
        # 测试 6: 清空历史记录
        test_clear_history(token)
        
        # 测试 7: 无效 Token
        test_invalid_token()
        
        # 总结
        print_section("✅ 测试完成")
        print("""
测试结果总结：
1. ✅ 用户注册功能正常
2. ✅ 用户登录返回 JWT Token
3. ✅ 无 Token 访问被拒绝（401）
4. ✅ 有效 Token 可以正常对话
5. ✅ 历史记录按用户隔离
6. ✅ 清空历史按用户隔离
7. ✅ 无效 Token 被拒绝

🎉 所有认证功能测试通过！
        """)
        
    except requests.exceptions.ConnectionError:
        print("\n❌ 连接失败！请确保服务器已启动：")
        print("   python server.py")
    except Exception as e:
        print(f"\n❌ 测试过程出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()