diff --git a/metagpt/actions/write_test.py b/metagpt/actions/write_test.py
index 393b5fcd9..73d4c119d 100644
--- a/metagpt/actions/write_test.py
+++ b/metagpt/actions/write_test.py
@@ -5,23 +5,54 @@
 @Author  : alexanderwu
 @File    : write_test.py
 """
+from metagpt.logs import logger
 from metagpt.actions.action import Action
+from metagpt.utils.common import CodeParser
 
+PROMPT_TEMPLATE = """
+NOTICE
+1. Role: You are a QA engineer; the main goal is to design, develop, and execute PEP8 compliant, well-structured, maintainable test cases and scripts for Python 3.9. Your focus should be on ensuring the product quality of the entire project through systematic testing.
+2. Requirement: Based on the context, develop a comprehensive test suite that adequately covers all relevant aspects of the code file under review. Your test suite will be part of the overall project QA, so please develop complete, robust, and reusable test cases.
+3. Attention1: Use '##' to split sections, not '#', and '## <SECTION_NAME>' SHOULD WRITE BEFORE the test case or script.
+4. Attention2: If there are any settings in your tests, ALWAYS SET A DEFAULT VALUE, ALWAYS USE STRONG TYPE AND EXPLICIT VARIABLE.
+5. Attention3: YOU MUST FOLLOW "Data structures and interface definitions". DO NOT CHANGE ANY DESIGN. Make sure your tests respect the existing design and ensure its validity.
+6. Think before writing: What should be tested and validated in this document? What edge cases could exist? What might fail?
+7. CAREFULLY CHECK THAT YOU DON'T MISS ANY NECESSARY TEST CASES/SCRIPTS IN THIS FILE.
+Attention: Use '##' to split sections, not '#', and '## <SECTION_NAME>' SHOULD WRITE BEFORE the test case or script and triple quotes.
+-----
+## Given the following code, please write appropriate test cases using Python's unittest framework to verify the correctness and robustness of this code:
+```python
+{code_to_test}
+```
+Note that the code to test is at {source_file_path}, we will put your test code at {workspace}/tests/{test_file_name}, and run your test code from {workspace},
+you should correctly import the necessary classes based on these file locations!
+## {test_file_name}: Write test code with triple quoto. Do your best to implement THIS ONLY ONE FILE.
+"""
 
 class WriteTest(Action):
-    def __init__(self, name="", context=None, llm=None):
+    def __init__(self, name="WriteTest", context=None, llm=None):
         super().__init__(name, context, llm)
-        self.code = None
-        self.test_prompt_template = """
-        Given the following code or function:
-        {code}
 
-        As a test engineer, please write appropriate test cases using Python's unittest framework to verify the correctness and robustness of this code.
-        """
+    async def write_code(self, prompt):
+        code_rsp = await self._aask(prompt)
+        code = CodeParser.parse_code(block="", text=code_rsp)
+        return code
 
+<<<<<<< main
     async def run(self, code):
         self.code = code
         prompt = self.test_prompt_template.format(code=self.code)
         test_cases = await self._aask(prompt)
         return test_cases
-    
\ No newline at end of file
+    
+=======
+    async def run(self, code_to_test, test_file_name, source_file_path, workspace):
+        prompt = PROMPT_TEMPLATE.format(
+            code_to_test=code_to_test,
+            test_file_name=test_file_name,
+            source_file_path=source_file_path,
+            workspace=workspace
+        )
+        code = await self.write_code(prompt)
+        return code
+>>>>>>> main