Agent/test_resource_storage.py

#!/usr/bin/env python3
"""测试Resource存储系统"""

import requests
import base64
import os

BASE_URL = "http://localhost:8000"

# 生成测试密钥
test_key = base64.b64encode(os.urandom(32)).decode('ascii')
print(f"测试密钥: {test_key}")
print(f"请在.env中设置: ORG_KEYS=test:{test_key}")
print()

def test_submit_resource():
    """测试提交resource"""
    print("=== 测试1: 提交普通代码 ===")
    response = requests.post(f"{BASE_URL}/api/resource", json={
        "id": "test/code/example",
        "title": "示例代码",
        "body": "print('Hello World')",
        "content_type": "code",
        "metadata": {"language": "python"},
        "submitted_by": "test@example.com"
    })
    print(f"状态码: {response.status_code}")
    print(f"响应: {response.json()}")
    print()

    print("=== 测试2: 提交敏感凭证 ===")
    response = requests.post(f"{BASE_URL}/api/resource", json={
        "id": "test/credentials/website",
        "title": "网站登录凭证",
        "body": "使用方法：直接登录",
        "secure_body": "账号：user@example.com\n密码：SecurePass123",
        "content_type": "credential",
        "metadata": {"acquired_at": "2026-03-06T10:00:00Z"},
        "submitted_by": "test@example.com"
    })
    print(f"状态码: {response.status_code}")
    print(f"响应: {response.json()}")
    print()

    print("=== 测试3: 提交Cookie ===")
    response = requests.post(f"{BASE_URL}/api/resource", json={
        "id": "test/cookies/website",
        "title": "网站Cookie",
        "body": "适用于：已登录状态",
        "secure_body": "session_id=abc123; auth_token=xyz789",
        "content_type": "cookie",
        "metadata": {
            "acquired_at": "2026-03-06T10:00:00Z",
            "expires_at": "2026-03-07T10:00:00Z"
        },
        "submitted_by": "test@example.com"
    })
    print(f"状态码: {response.status_code}")
    print(f"响应: {response.json()}")
    print()


def test_get_resource():
    """测试获取resource"""
    print("=== 测试4: 获取普通代码（无需密钥）===")
    response = requests.get(f"{BASE_URL}/api/resource/test/code/example")
    print(f"状态码: {response.status_code}")
    data = response.json()
    print(f"标题: {data['title']}")
    print(f"内容: {data['body']}")
    print(f"敏感内容: {data['secure_body']}")
    print()

    print("=== 测试5: 获取凭证（无密钥）===")
    response = requests.get(f"{BASE_URL}/api/resource/test/credentials/website")
    print(f"状态码: {response.status_code}")
    data = response.json()
    print(f"标题: {data['title']}")
    print(f"公开内容: {data['body']}")
    print(f"敏感内容: {data['secure_body']}")  # 应该是 [ENCRYPTED]
    print()

    print("=== 测试6: 获取凭证（有密钥）===")
    response = requests.get(
        f"{BASE_URL}/api/resource/test/credentials/website",
        headers={"X-Org-Key": test_key}
    )
    print(f"状态码: {response.status_code}")
    data = response.json()
    print(f"标题: {data['title']}")
    print(f"公开内容: {data['body']}")
    print(f"敏感内容: {data['secure_body']}")  # 应该解密
    print()


def test_patch_resource():
    """测试更新resource"""
    print("=== 测试7: 更新resource ===")
    response = requests.patch(
        f"{BASE_URL}/api/resource/test/credentials/website",
        json={
            "title": "更新后的标题",
            "metadata": {"acquired_at": "2026-03-06T11:00:00Z"}
        }
    )
    print(f"状态码: {response.status_code}")
    print(f"响应: {response.json()}")
    print()


def test_list_resources():
    """测试列出resource"""
    print("=== 测试8: 列出所有resource ===")
    response = requests.get(f"{BASE_URL}/api/resource")
    print(f"状态码: {response.status_code}")
    data = response.json()
    print(f"总数: {data['count']}")
    for item in data['results']:
        print(f"  - {item['id']}: {item['title']} ({item['content_type']})")
    print()

    print("=== 测试9: 按类型过滤 ===")
    response = requests.get(f"{BASE_URL}/api/resource?content_type=credential")
    print(f"状态码: {response.status_code}")
    data = response.json()
    print(f"凭证数量: {data['count']}")
    for item in data['results']:
        print(f"  - {item['id']}: {item['title']}")
    print()


if __name__ == "__main__":
    print("请确保服务器正在运行: uvicorn knowhub.server:app --reload")
    print()

    try:
        test_submit_resource()
        test_get_resource()
        test_patch_resource()
        test_list_resources()
        print("✅ 所有测试完成")
    except Exception as e:
        print(f"❌ 测试失败: {e}")