{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Control Flow Graph (CFG) Representation for the `generateBoard()` method in `EightQueens` class\n",
      "\n",
      "#### Nodes:\n",
      "1. `Queen[] start = new Queen[8];` - Declaration and initialization of Queen array.\n",
      "2. `Random gen = new Random();` - Initialization of Random generator object.\n",
      "3. `int i = 0;` - Initialization of loop variable.\n",
      "4. `i < 8;` - Loop condition check.\n",
      "5. `start[i] = new Queen(gen.nextInt(8), i);` - Assigning new Queen object to array element during each iteration.\n",
      "6. `i++;` - Increment of loop variable.\n",
      "7. `return start;` - Return the array of Queen objects.\n",
      "8. `End of generateBoard()` - Exit point of the method.\n",
      "\n",
      "#### Edges:\n",
      "- `Edge 1`: Node 1 to Node 2 - Sequence of operations.\n",
      "- `Edge 2`: Node 2 to Node 3 - Sequence transitioning to loop initialization.\n",
      "- `Edge 3`: Node 3 to Node 4 - Moves from loop initialization to loop condition check.\n",
      "- `Edge 4`: Node 4 to Node 5 (Condition: true) - If condition is true, execute loop body.\n",
      "- `Edge 5`: Node 5 to Node 6 - Sequence inside loop, after execution moves to increment statement.\n",
      "- `Edge 6`: Node 6 to Node 4 - After incrementing, control goes back to loop condition check.\n",
      "- `Edge 7`: Node 4 to Node 7 (Condition: false) - If loop condition fails, flow exits loop and returns array.\n",
      "- `Edge 8`: Node 7 to Node 8 - Control exits the method.\n",
      "\n",
      "#### Conditions:\n",
      "- Node 4 (`i < 8;`): Determines whether the loop should continue or exit.\n",
      "\n",
      "#### Loops:\n",
      "- **Loop Start**: Node 3 (`int i = 0;`)\n",
      "- **Loop End**: Node 7 (`return start;`)\n",
      "- **Condition**: `i < 8;` - Ensures loop iterates through the indices of the array.\n",
      "\n",
      "This CFG allows tracing the execution flow in the `generateBoard()` method, helping to understand the iterative process that populates the `Queen[]` array with newly created `Queen` objects, where each object is placed randomly on the board while its index maintains the column placement. This CFG is particularly useful for debugging or extending the method due to its clear representation of each computational step and control transfer.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import glob\n",
    "import re\n",
    "import os\n",
    "from openai import OpenAI\n",
    "\n",
    "api_key = 'sk-ZLMKmXkh5EAlsGE9431f6c690cE6498c847f48C7D7FbAfA6'\n",
    "api_base = \"https://openkey.cloud/v1\"\n",
    "client = OpenAI(api_key=api_key, base_url=api_base)\n",
    "test_methods = \"\"\"\n",
    "package net.mooctest;\n",
    "/**\n",
    " * @file EightQueens.java\n",
    " * @author nsquires\n",
    " * Solves the eight queens problem using various AI techniques\n",
    " */\n",
    "import java.util.*;\n",
    "import java.text.NumberFormat;\n",
    "\n",
    "public class EightQueens {\n",
    "\t\n",
    "\tpublic EightQueens(){\n",
    "\t}\n",
    "\t\n",
    "\t/**\n",
    "\t * The starter board\n",
    "\t * @return Queen[]\n",
    "\t */\n",
    "\tpublic Queen[] generateBoard(){\n",
    "\t\tQueen[] start = new Queen[8];\n",
    "\t\tRandom gen = new Random();\n",
    "\t\t\n",
    "\t\tfor(int i=0; i<8; i++){\n",
    "\t\t\tstart[i] = new Queen(gen.nextInt(8),i);\n",
    "\t\t}\n",
    "\t\treturn start;\n",
    "\t}\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "def getCFG(java_code):     \n",
    "    prompt = f\"\"\"\n",
    "Please generate a Control Flow Graph (CFG) from the following Java code. The CFG should include detailed information about:\n",
    "\n",
    "1. Nodes: Identify each independent statement or expression as a node.\n",
    "   - Purpose: Nodes represent points of computation or decision in the code.\n",
    "\n",
    "2. Edges: Describe the control flow transitions between nodes.\n",
    "   - Purpose: Edges show how execution moves from one node to another, including conditions that influence the flow.\n",
    "\n",
    "3. Conditions: Specify the conditions for decision-making statements (if, switch, loops).\n",
    "   - Purpose: Conditions determine the branching in the flow of execution.\n",
    "\n",
    "4. Loops: Identify the start and end of loops, and describe their iteration conditions.\n",
    "   - Purpose: Loops are critical for understanding repeated execution and potential paths within the code.\n",
    "\n",
    "Please format the CFG as follows:\n",
    "- Nodes: List all nodes with a brief description.\n",
    "- Edges: List all edges with starting and ending node identifiers and the condition (if applicable).\n",
    "- Conditions: List all conditions associated with their nodes.\n",
    "- Loops: Describe all loops with their starting node, ending node, and condition.\n",
    "\n",
    "Java Code:\n",
    "{java_code}\n",
    "\n",
    "\"\"\"\n",
    "    response = client.chat.completions.create(\n",
    "    model=\"gpt-3.5-turbo\",\n",
    "#   response_format={ \"type\": \"json_object\" },\n",
    "    messages=[\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "    ]\n",
    "    )\n",
    "    print(response.choices[0].message.content)\n",
    "\n",
    "getCFG(test_methods)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['D:\\\\learn\\\\junit-test\\\\test\\\\EightQueens_1665290554002\\\\src\\\\test\\\\java\\\\net\\\\mooctest\\\\EightQueensTest.java']\n",
      "find method count 1\n",
      "find method: @Test(timeout = 4000)\n",
      "\tpublic void test()\n",
      "find: test @Test(timeout = 4000)\n",
      "\tpublic void test() {\n",
      "\t\tEightQueens board = new EightQueens();\n",
      "\t\tQueen[] startBoard = board.generateBoard();\n",
      "\t\t// 检查不攻击的女王触发\n",
      "\t\tassertTrue(new Queen(1, 3).canAttack(new Queen(2, 2)));\n",
      "\t\t// 检查同一行的女王触发\n",
      "\t\tassertTrue(new Queen(1, 3).canAttack(new Queen(1, 5)));\n",
      "\t\t// 检查同一列的女王触发\n",
      "\t\tassertTrue(new Queen(3, 1).canAttack(new Queen(5, 1)));\n",
      "\t\t// 检查同一对角线的女王触发\n",
      "\t\tassertTrue(new Queen(1, 1).canAttack(new Queen(2, 2)));\n",
      "\t\t// 2. movedown\n",
      "\t\tstartBoard[0].setRow(0);\n",
      "\t\tstartBoard[0].setColumn(1);\n",
      "\t\tassertEquals(startBoard[0].getRow(), 0);\n",
      "\t\tassertEquals(startBoard[0].getColumn(), 1);\n",
      "\t\tfor(int i = 0; i < 30; i++) startBoard[0].moveDown(i);\n",
      "\t\t// 3. Test HillClimbing  爬山法解八皇后问题\n",
      "\t\tHillClimbing hcTest = new HillClimbing();\n",
      "\t\tHillClimbing hc = new HillClimbing(startBoard);\n",
      "\t\t// 4. 检查生成的Node\n",
      "\t\tNode node = hc.hillClimbing();\n",
      "\t\tNode nodeTest = hcTest.getStartNode();\n",
      "\t\tassertNotNull(node);\n",
      "\t\tassertNotNull(nodeTest);\n",
      "\t\tassertNotNull(node.getRandomNeighbour(node));\n",
      "\t\t// 5. 测试RandomRestart\n",
      "\t\tRandomRestart rr = new RandomRestart(startBoard);\n",
      "\t\tNode randomNode = rr.randomRestart();\n",
      "\t\tassertNotNull(rr);\n",
      "\t\tassertNotNull(randomNode);\n",
      "\t\tNode now = rr.randomRestart();\n",
      "\t\tassertNotNull(now);\n",
      "\t\trr.setStartNode(now);\n",
      "\t\tassertSame(rr.getStartNode(), now);\n",
      "\t\t// 6. 测试SimulatedAnnealing    模拟退火解八皇后问题\n",
      "\t\tSimulatedAnnealing sa = new SimulatedAnnealing(startBoard);\n",
      "\t\tsa.startState();\n",
      "\t\tassertNotNull(sa.getStartNode());\n",
      "\t\tNode ans;\n",
      "\t\tdouble T = 1000.0, delta = 0.7;    // 温度越高越好，速率0.8~0.99\n",
      "\t\twhile(delta < 1) {\n",
      "\t\t\tans = sa.simulatedAnneal(T, delta);\n",
      "\t\t\tassertNotNull(ans);\n",
      "\t\t\tdelta *= 1.1;\n",
      "\t\t\tT *= 1.1;\n",
      "\t\t}\n",
      "\t\t\n",
      "\t}\n",
      "Method 'test' was successfully split into separate test units.\n",
      "be replace test\n",
      "old------------- @Test(timeout = 4000)\n",
      "\tpublic void test() {\n",
      "\t\tEightQueens board = new EightQueens();\n",
      "\t\tQueen[] startBoard = board.generateBoard();\n",
      "\t\t// 检查不攻击的女王触发\n",
      "\t\tassertTrue(new Queen(1, 3).canAttack(new Queen(2, 2)));\n",
      "\t\t// 检查同一行的女王触发\n",
      "\t\tassertTrue(new Queen(1, 3).canAttack(new Queen(1, 5)));\n",
      "\t\t// 检查同一列的女王触发\n",
      "\t\tassertTrue(new Queen(3, 1).canAttack(new Queen(5, 1)));\n",
      "\t\t// 检查同一对角线的女王触发\n",
      "\t\tassertTrue(new Queen(1, 1).canAttack(new Queen(2, 2)));\n",
      "\t\t// 2. movedown\n",
      "\t\tstartBoard[0].setRow(0);\n",
      "\t\tstartBoard[0].setColumn(1);\n",
      "\t\tassertEquals(startBoard[0].getRow(), 0);\n",
      "\t\tassertEquals(startBoard[0].getColumn(), 1);\n",
      "\t\tfor(int i = 0; i < 30; i++) startBoard[0].moveDown(i);\n",
      "\t\t// 3. Test HillClimbing  爬山法解八皇后问题\n",
      "\t\tHillClimbing hcTest = new HillClimbing();\n",
      "\t\tHillClimbing hc = new HillClimbing(startBoard);\n",
      "\t\t// 4. 检查生成的Node\n",
      "\t\tNode node = hc.hillClimbing();\n",
      "\t\tNode nodeTest = hcTest.getStartNode();\n",
      "\t\tassertNotNull(node);\n",
      "\t\tassertNotNull(nodeTest);\n",
      "\t\tassertNotNull(node.getRandomNeighbour(node));\n",
      "\t\t// 5. 测试RandomRestart\n",
      "\t\tRandomRestart rr = new RandomRestart(startBoard);\n",
      "\t\tNode randomNode = rr.randomRestart();\n",
      "\t\tassertNotNull(rr);\n",
      "\t\tassertNotNull(randomNode);\n",
      "\t\tNode now = rr.randomRestart();\n",
      "\t\tassertNotNull(now);\n",
      "\t\trr.setStartNode(now);\n",
      "\t\tassertSame(rr.getStartNode(), now);\n",
      "\t\t// 6. 测试SimulatedAnnealing    模拟退火解八皇后问题\n",
      "\t\tSimulatedAnnealing sa = new SimulatedAnnealing(startBoard);\n",
      "\t\tsa.startState();\n",
      "\t\tassertNotNull(sa.getStartNode());\n",
      "\t\tNode ans;\n",
      "\t\tdouble T = 1000.0, delta = 0.7;    // 温度越高越好，速率0.8~0.99\n",
      "\t\twhile(delta < 1) {\n",
      "\t\t\tans = sa.simulatedAnneal(T, delta);\n",
      "\t\t\tassertNotNull(ans);\n",
      "\t\t\tdelta *= 1.1;\n",
      "\t\t\tT *= 1.1;\n",
      "\t\t}\n",
      "\t\t\n",
      "\t}\n",
      "new------------- @Test(timeout = 4000)\n",
      "\tpublic void testGenerateBoard() {\n",
      "\t\tEightQueens board = new EightQueens();\n",
      "\t\tQueen[] startBoard = board.generateBoard();\n",
      "\t\tassertNotNull(startBoard);\n",
      "\t}\n",
      "\n",
      "\t@Test\n",
      "\tpublic void testCanAttack() {\n",
      "\t\tassertTrue(new Queen(1, 3).canAttack(new Queen(2, 2)));\n",
      "\t\tassertTrue(new Queen(1, 3).canAttack(new Queen(1, 5)));\n",
      "\t\tassertTrue(new Queen(3, 1).canAttack(new Queen(5, 1)));\n",
      "\t\tassertTrue(new Queen(1, 1).canAttack(new Queen(2, 2)));\n",
      "\t}\n",
      "\n",
      "\t@Test\n",
      "\tpublic void testMoveDown() {\n",
      "\t\tQueen queen = new Queen(1, 3);\n",
      "\t\tqueen.moveDown(5);\n",
      "\t\tassertEquals(queen.getRow(), 6);\n",
      "\t\tassertEquals(queen.getColumn(), 3);\n",
      "\t}\n",
      "\n",
      "\t@Test\n",
      "\tpublic void testHillClimbing() {\n",
      "\t\tHillClimbing hc = new HillClimbing();\n",
      "\t\thc.startState();\n",
      "\t\tNode node = hc.hillClimbing();\n",
      "\t\tassertNotNull(node);\n",
      "\t}\n",
      "\n",
      "\t@Test\n",
      "\tpublic void testRandomRestart() {\n",
      "\t\tEightQueens board = new EightQueens();\n",
      "\t\tQueen[] startBoard = board.generateBoard();\n",
      "\t\tRandomRestart rr = new RandomRestart(startBoard);\n",
      "\t\tNode node = rr.randomRestart();\n",
      "\t\tassertNotNull(node);\n",
      "\t\tNode startNode = rr.getStartNode();\n",
      "\t\trr.setStartNode(startNode);\n",
      "\t\tassertSame(rr.getStartNode(), startNode);\n",
      "\t}\n",
      "\n",
      "\t@Test\n",
      "\tpublic void testSimulatedAnnealing() {\n",
      "\t\tEightQueens board = new EightQueens();\n",
      "\t\tQueen[] startBoard = board.generateBoard();\n",
      "\t\tSimulatedAnnealing sa = new SimulatedAnnealing(startBoard);\n",
      "\t\tsa.startState();\n",
      "\t\tNode startNode = sa.getStartNode();\n",
      "\t\tassertNotNull(startNode);\n",
      "\t\tdouble initialTemp = 1000.0;\n",
      "\t\tdouble step = 0.7;\n",
      "\t\tNode ans = sa.simulatedAnneal(initialTemp, step);\n",
      "\t\tassertNotNull(ans);\n",
      "\t}\n"
     ]
    }
   ],
   "source": [
    "import glob\n",
    "import re\n",
    "import os\n",
    "from openai import OpenAI\n",
    "\n",
    "api_key = 'sk-ZLMKmXkh5EAlsGE9431f6c690cE6498c847f48C7D7FbAfA6'\n",
    "api_base = \"https://openkey.cloud/v1\"\n",
    "client = OpenAI(api_key=api_key, base_url=api_base)\n",
    "\n",
    "def get_test_methods(file_path):\n",
    "    with open(file_path, 'r', encoding='utf-8') as file:\n",
    "        content = file.read()\n",
    "    \n",
    "    method_declarations = re.findall(r'(@Test\\s*(?:\\(.*?\\))?\\s*public void (\\w+)\\(\\))', content)\n",
    "\n",
    "    test_methods = []\n",
    "    print(\"find method count\", len(method_declarations))\n",
    "    for method_declaration in method_declarations:\n",
    "        method_body = method_declaration[0]\n",
    "        print(\"find method:\", method_body)\n",
    "        origin_index = content.index(method_body)\n",
    "        start_index = origin_index + len(method_body)\n",
    "        # 继续向前查找，直到找到第一个 {\n",
    "        while content[start_index] != '{':\n",
    "            start_index += 1\n",
    "        brace_count = 1  # 已经找到了一个 {\n",
    "        for i in range(start_index + 1, len(content)):\n",
    "            if content[i] == '{':\n",
    "                brace_count += 1\n",
    "            elif content[i] == '}':\n",
    "                brace_count -= 1\n",
    "            if brace_count == 0:\n",
    "                # 找到了完整的方法体\n",
    "                method_body = content[origin_index:i+1]\n",
    "                test_methods.append((method_body, method_declaration[1]))\n",
    "                break\n",
    "\n",
    "# 移除不包含断言的方法\n",
    "    filtered_test_methods = [(method_name, method_body) for method_body, method_name in test_methods if re.search(r'\\b(?:assert\\w*|assertTrue|assertFalse)\\s*\\(', method_body)]\n",
    "    for (name, body) in filtered_test_methods:\n",
    "        print(\"find:\", name, body)\n",
    "    return filtered_test_methods\n",
    "\n",
    "def split_assertions(test_methods):\n",
    "    new_test_methods = []\n",
    "    for method_name, method_body in test_methods:\n",
    "        # 使用OpenAI对方法体进行拆分\n",
    "        try:\n",
    "            response = client.chat.completions.create(\n",
    "                model=\"gpt-3.5-turbo\",\n",
    "                messages=[\n",
    "            {\"role\": \"user\",\n",
    "            \"content\": f\"\"\"\n",
    "        你是一个拆分代码的AI助手，你的任务是将一个junit代码方法拆分成若干个独立可运行的junit测试单元方法。要求：\n",
    "        1.一个测试单元只能有一个断言\n",
    "        2.如果原代码只有一个断言，不用拆分\n",
    "        3.你的输入是一个完整的方法\n",
    "        4.输出该方法被拆分后的若干个方法\n",
    "        5.输出不需要其他提示词\n",
    "        方法体:\n",
    "        \\n\\n{method_body}\\n\\n\"\"\"}\n",
    "        ],\n",
    "                # response_format={ \"type\": \"json_object\" }\n",
    "            )\n",
    "            new_body = response.choices[0].message.content\n",
    "            new_test_methods.append((method_name, new_body))\n",
    "            print(f\"Method '{method_name}' was successfully split into separate test units.\")\n",
    "        except Exception as e:\n",
    "            print(f\"Failed to split method '{method_name}'. Error: {str(e)}\")\n",
    "    return new_test_methods\n",
    "\n",
    "def write_new_file(file_path, origin_test_methods, new_test_methods):\n",
    "    with open(file_path, 'r', encoding='utf-8') as file:\n",
    "        content = file.read()\n",
    "\n",
    "    # 创建一个字典，将方法名映射到新的方法体\n",
    "    new_methods_dict = {name: body for name, body in new_test_methods}\n",
    "\n",
    "    # 替换原始文件中的测试方法体\n",
    "    for method_name, old_body in origin_test_methods:\n",
    "        if method_name in new_methods_dict:\n",
    "            print(\"be replace\", method_name)\n",
    "            new_body = new_methods_dict[method_name]\n",
    "            print(\"old-------------\", old_body)\n",
    "            print(\"new-------------\", new_body)\n",
    "            content = content.replace(old_body, new_body)\n",
    "\n",
    "    # 将更改后的内容写回文件\n",
    "    with open(file_path, 'w', encoding='utf-8') as file:\n",
    "        file.write(content)\n",
    "\n",
    "def process_directory(directory):\n",
    "    test_files = glob.glob(os.path.join(directory, '**/src/test/**/*.java'), recursive=True)\n",
    "    print(test_files)\n",
    "    for file_path in test_files:\n",
    "        test_methods = get_test_methods(file_path)\n",
    "        if test_methods:\n",
    "            new_test_methods = split_assertions(test_methods)\n",
    "            write_new_file(file_path, test_methods, new_test_methods)\n",
    "\n",
    "# 指定要扫描的目录\n",
    "directory_to_scan = 'D:\\\\learn\\junit-test\\\\test\\\\EightQueens_1665290554002'\n",
    "process_directory(directory_to_scan)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}