From c247a8f1dc67889a8c71be1b1c5958cbde2695d2 Mon Sep 17 00:00:00 2001
From: sontolau <sonto.lau@gmail.com>
Date: Mon, 9 Mar 2026 10:47:52 +0800
Subject: [PATCH] source optimization

---
 ai_test_generator/config.py              |  1 +
 ai_test_generator/core/llm_client.py     |  2 +-
 ai_test_generator/core/parser.py         | 17 ++++++++++++-----
 ai_test_generator/core/prompt_builder.py |  4 ++--
 ai_test_generator/core/test_generator.py |  2 +-
 ai_test_generator/core/test_runner.py    |  2 +-
 ai_test_generator/run.sh                 |  4 ++--
 7 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/ai_test_generator/config.py b/ai_test_generator/config.py
index 6d416ac..563f7a6 100644
--- a/ai_test_generator/config.py
+++ b/ai_test_generator/config.py
@@ -17,5 +17,6 @@ class Config:
     # HTTP测试配置
     HTTP_BASE_URL: str = os.getenv("HTTP_BASE_URL", "http://localhost:8080")
     HTTP_TIMEOUT: int = 10
+    UNIT_KEYWORDS: tuple = ("units", "functions", "apis")
 
 config = Config()
\ No newline at end of file
diff --git a/ai_test_generator/core/llm_client.py b/ai_test_generator/core/llm_client.py
index 312ec59..a74f090 100644
--- a/ai_test_generator/core/llm_client.py
+++ b/ai_test_generator/core/llm_client.py
@@ -2,7 +2,7 @@ import json
 import re
 import logging
 from openai import OpenAI
-from gui_ai_developer.config import config
+from config import config
 
 logger = logging.getLogger(__name__)
 
diff --git a/ai_test_generator/core/parser.py b/ai_test_generator/core/parser.py
index b90acaa..dbe5b23 100644
--- a/ai_test_generator/core/parser.py
+++ b/ai_test_generator/core/parser.py
@@ -32,7 +32,7 @@ import json
 import re
 from typing import Any
 from dataclasses import dataclass, field
-
+from config import config
 
 # ══════════════════════════════════════════════════════════════
 # 基础数据结构
@@ -213,11 +213,18 @@ class InterfaceParser:
                 description="",
                 interfaces=self._parse_units(data),
             )
+
+        units = []
+        for unit in config.UNIT_KEYWORDS:
+            if unit not in data:
+                continue
+            units = data.get(unit, [])
+
         return ApiDescriptor(
-            project=data.get("project", "default"),
-            description=data.get("description", ""),
-            interfaces=self._parse_units(data.get("units", [])),
-        )
+                project=data.get("project", "default"),
+                description=data.get("description", ""),
+                interfaces=self._parse_units(units),
+            )
 
     # ── 单元解析 ──────────────────────────────────────────────
 
diff --git a/ai_test_generator/core/prompt_builder.py b/ai_test_generator/core/prompt_builder.py
index 14e82d4..1612df2 100644
--- a/ai_test_generator/core/prompt_builder.py
+++ b/ai_test_generator/core/prompt_builder.py
@@ -155,8 +155,8 @@ class PromptBuilder:
             f"## Available Interfaces\n"
             f"{json.dumps(iface_summary, ensure_ascii=False, indent=2)}\n\n"
             f"## Test Requirements\n"
-            f"{req_lines}\n\n"
-            f"Generate comprehensive test cases (positive + negative) for every requirement.\n"
+            f"Generate comprehensive test cases (positive + negative) for every interface.\n"
+            f"- The generated test cases for each interface must meet the following requirements: {req_lines}\n"
             f"- For function interfaces: use 'module_path' for import, "
             f"'source_file' for reference.\n"
             f"- For HTTP interfaces: use 'full_url' as the request URL.\n"
diff --git a/ai_test_generator/core/test_generator.py b/ai_test_generator/core/test_generator.py
index be432d1..ebe4365 100644
--- a/ai_test_generator/core/test_generator.py
+++ b/ai_test_generator/core/test_generator.py
@@ -11,7 +11,7 @@ import json
 import re
 import logging
 from pathlib import Path
-from gui_ai_developer.config import config
+from config import config
 
 logger = logging.getLogger(__name__)
 
diff --git a/ai_test_generator/core/test_runner.py b/ai_test_generator/core/test_runner.py
index 5bd8f7c..90a3e7f 100644
--- a/ai_test_generator/core/test_runner.py
+++ b/ai_test_generator/core/test_runner.py
@@ -6,7 +6,7 @@ import time
 import logging
 from pathlib import Path
 from dataclasses import dataclass, field
-from gui_ai_developer.config import config
+from config import config
 
 logger = logging.getLogger(__name__)
 
diff --git a/ai_test_generator/run.sh b/ai_test_generator/run.sh
index 4cc0dc3..6da9f8c 100755
--- a/ai_test_generator/run.sh
+++ b/ai_test_generator/run.sh
@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 
 # 安装依赖
-pip install -r requirements.txt
+# pip install -r requirements.txt
 
 # 设置环境变量
 export LLM_API_KEY="sk-AUmOuFI731Ty5Nob38jY26d8lydfDT-QkE2giqb0sCuPCAE2JH6zjLM4lZLpvL5WMYPOocaMe2FwVDmqM_9KimmKACjR"
@@ -15,5 +15,5 @@ export HTTP_BASE_URL="http://localhost:8080"
 #  --requirements "创建用户，自动生成测试数据,更改指定用户的密码，自动生成测试数据"
 
 # 只生成不执行
-python main.py --api-desc examples/function_signatures.json \
+python main.py --api-desc examples/api_desc.json \
   --requirements "1.对每个接口进行测试，支持特殊值、边界值测试"