codeAnalysis/test_gen_demo.py at master · YMX899/codeAnalysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import requests
import json
import sys

BASE_URL = "http://127.0.0.1:8000/api"

def print_json(data):
    print(json.dumps(data, indent=2, ensure_ascii=False))

def check_health():
    try:
        resp = requests.get("http://127.0.0.1:8000/health")
        if resp.status_code == 200:
            print("✅ 服务健康检查通过")
            return True
        else:
            print(f"❌ 服务健康检查失败: {resp.status_code}")
            return False
    except Exception as e:
        print(f"❌ 无法连接服务: {e}")
        print("请确保已运行 python run_server.py")
        return False

def get_first_project():
    resp = requests.get(f"{BASE_URL}/analysis/projects")
    if resp.status_code != 200:
        print(f"❌ 获取项目列表失败: {resp.text}")
        return None

    data = resp.json()
    projects = data.get("projects", [])
    if not projects:
        print("⚠️ 数据库中没有已分析的项目。")
        print("💡 请先使用 POST /api/analysis/analyze 接口分析一个 Java 项目。")
        return None

    project = projects[0]
    print(f"✅ 找到项目: {project['name']} (ID: {project['id']})")
    return project

def get_random_method(project_id):
    # 搜索所有方法 (空关键词可能搜不到，试个常用词 "get" 或 "main" 或空字符串如果支持)
    # GraphRepository 的 search_methods 实现是用 CONTAINS $keyword，空字符串应该匹配所有
    search_payload = {
        "keyword": "",  # 空字符串尝试匹配所有
        "element_type": "method",
        "limit": 5
    }

    # 这里的 search 接口是 POST /search/{project_id}
    resp = requests.post(f"{BASE_URL}/analysis/search/{project_id}", json=search_payload)

    if resp.status_code != 200:
        print(f"❌ 搜索方法失败: {resp.text}")
        return None

    results = resp.json().get("results", [])
    if not results:
        # 尝试搜 "get"
        search_payload["keyword"] = "get"
        resp = requests.post(f"{BASE_URL}/analysis/search/{project_id}", json=search_payload)
        results = resp.json().get("results", [])

    if not results:
        print("⚠️ 未能在项目中找到任何方法。")
        return None

    method = results[0]
    print(f"✅ 找到目标方法: {method['method_name']} (ID: {method['method_id']})")
    print(f"   签名: {method['signature']}")
    print(f"   所属类: {method.get('class_name', 'Unknown')}")
    return method

def generate_test(project_id, method_id):
    payload = {
        "project_id": project_id,
        "target_element_id": method_id,
        "element_type": "method",
        "test_framework": "junit",
        "additional_instructions": "请添加详细的代码注释，并使用 Mockito 模拟外部依赖。"
    }

    print("\n⏳ 正在请求 AI 生成测试代码 (这可能需要几秒钟)...")
    resp = requests.post(f"{BASE_URL}/tests/generate", json=payload)

    if resp.status_code != 200:
        print(f"❌ 生成测试失败: {resp.text}")
        return

    data = resp.json()
    print("\n🎉 测试生成成功！")
    print(f"📂 建议文件名: {data['suggested_filename']}")
    print("-" * 50)
    print(data['test_code'])
    print("-" * 50)

def main():
    if not check_health():
        return

    project = get_first_project()
    if not project:
        return

    method = get_random_method(project['id'])
    if not method:
        return

    generate_test(project['id'], method['method_id'])

if __name__ == "__main__":
    main()