diff --git a/unit_tests/locales/base.pot b/unit_tests/locales/base.pot
index 7752859..e560b2b 100644
--- a/unit_tests/locales/base.pot
+++ b/unit_tests/locales/base.pot
@@ -8,7 +8,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-01-11 21:35+0800\n"
+"POT-Creation-Date: 2024-01-18 15:29+0800\n"
 "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n"
 "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
 "Language-Team: LANGUAGE <LL@li.org>\n"
@@ -17,44 +17,65 @@ msgstr ""
 "Content-Type: text/plain; charset=CHARSET\n"
 "Content-Transfer-Encoding: 8bit\n"
 
-#: main.py:30
+#: main.py:53
 msgid "Analyzing the function and current unit tests..."
 msgstr ""
 
-#: main.py:49
+#: main.py:84
 msgid "Select test cases to generate"
 msgstr ""
 
-#: main.py:54
+#: main.py:89
+msgid ""
+"You can add more test cases here\n"
+"(Multiple cases can be separated by line breaks)"
+msgstr ""
+
+#: main.py:95
 msgid ""
 "Edit reference test file\n"
 "(Multiple files can be separated by line breaks)"
 msgstr ""
 
-#: main.py:68
+#: main.py:112
 msgid "No test case is selected. Quit generating tests."
 msgstr ""
 
-#: main.py:90
-msgid "No valid file is provided. Will not use reference to generate tests."
+#: main.py:129
+msgid "Will generate tests for the following cases."
 msgstr ""
 
-#: main.py:93
-msgid "Will use the following reference files to generate tests:"
+#: main.py:132
+msgid ""
+"\n"
+"Test cases:"
 msgstr ""
 
-#: main.py:94
+#: main.py:139
+msgid ""
+"\n"
+"No valid reference file is provided. Will not use reference to generate "
+"tests."
+msgstr ""
+
+#: main.py:144
+msgid ""
+"\n"
+"Will use the following reference files to generate tests."
+msgstr ""
+
+#: main.py:145
 msgid ""
 "\n"
 "Valid reference files:"
 msgstr ""
 
-#: main.py:98
+#: main.py:150
 msgid ""
 "\n"
 "Invalid files:"
 msgstr ""
 
-#: main.py:169
+#: main.py:232
 msgid "The function's size surpasses AI's context capacity."
 msgstr ""
diff --git a/unit_tests/locales/en/base.po b/unit_tests/locales/en/base.po
index bfa5eb1..cfbd0e4 100644
--- a/unit_tests/locales/en/base.po
+++ b/unit_tests/locales/en/base.po
@@ -7,7 +7,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-01-11 21:35+0800\n"
+"POT-Creation-Date: 2024-01-18 15:29+0800\n"
 "PO-Revision-Date: 2023-12-24 16:51+0800\n"
 "Last-Translator: kagami <mingjing@merico.dev>\n"
 "Language-Team: English\n"
@@ -17,45 +17,66 @@ msgstr ""
 "Content-Transfer-Encoding: 8bit\n"
 "Plural-Forms: nplurals=2; plural=(n != 1);\n"
 
-#: main.py:30
+#: main.py:53
 msgid "Analyzing the function and current unit tests..."
 msgstr ""
 
-#: main.py:49
+#: main.py:84
 msgid "Select test cases to generate"
 msgstr ""
 
-#: main.py:54
+#: main.py:89
+msgid ""
+"You can add more test cases here\n"
+"(Multiple cases can be separated by line breaks)"
+msgstr ""
+
+#: main.py:95
 msgid ""
 "Edit reference test file\n"
 "(Multiple files can be separated by line breaks)"
 msgstr ""
 
-#: main.py:68
+#: main.py:112
 msgid "No test case is selected. Quit generating tests."
 msgstr ""
 
-#: main.py:90
-msgid "No valid file is provided. Will not use reference to generate tests."
+#: main.py:129
+msgid "Will generate tests for the following cases."
 msgstr ""
 
-#: main.py:93
-msgid "Will use the following reference files to generate tests:"
+#: main.py:132
+msgid ""
+"\n"
+"Test cases:"
 msgstr ""
 
-#: main.py:94
+#: main.py:139
+msgid ""
+"\n"
+"No valid reference file is provided. Will not use reference to generate "
+"tests."
+msgstr ""
+
+#: main.py:144
+msgid ""
+"\n"
+"Will use the following reference files to generate tests."
+msgstr ""
+
+#: main.py:145
 msgid ""
 "\n"
 "Valid reference files:"
 msgstr ""
 
-#: main.py:98
+#: main.py:150
 msgid ""
 "\n"
 "Invalid files:"
 msgstr ""
 
-#: main.py:169
+#: main.py:232
 msgid "The function's size surpasses AI's context capacity."
 msgstr ""
 
diff --git a/unit_tests/locales/zh/LC_MESSAGES/base.mo b/unit_tests/locales/zh/LC_MESSAGES/base.mo
index acd52f1..3cbdf33 100644
Binary files a/unit_tests/locales/zh/LC_MESSAGES/base.mo and b/unit_tests/locales/zh/LC_MESSAGES/base.mo differ
diff --git a/unit_tests/locales/zh/base.po b/unit_tests/locales/zh/base.po
index 6256c4f..be1ea7e 100644
--- a/unit_tests/locales/zh/base.po
+++ b/unit_tests/locales/zh/base.po
@@ -7,7 +7,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-01-11 21:35+0800\n"
+"POT-Creation-Date: 2024-01-18 15:29+0800\n"
 "PO-Revision-Date: 2023-12-24 16:51+0800\n"
 "Last-Translator: kagami <mingjing@merico.dev>\n"
 "Language-Team: Chinese\n"
@@ -16,15 +16,23 @@ msgstr ""
 "Content-Type: text/plain; charset=UTF-8\n"
 "Content-Transfer-Encoding: 8bit\n"
 
-#: main.py:30
+#: main.py:53
 msgid "Analyzing the function and current unit tests..."
 msgstr "正在分析目标函数和项目中现有的单元测试……"
 
-#: main.py:49
+#: main.py:84
 msgid "Select test cases to generate"
 msgstr "选择测试用例以生成单元测试"
 
-#: main.py:54
+#: main.py:89
+msgid ""
+"You can add more test cases here\n"
+"(Multiple cases can be separated by line breaks)"
+msgstr ""
+"可在输入框中添加更多测试用例\n"
+"（多个测试用例用换行分隔）"
+
+#: main.py:95
 msgid ""
 "Edit reference test file\n"
 "(Multiple files can be separated by line breaks)"
@@ -32,19 +40,34 @@ msgstr ""
 "参考的测试文件\n"
 "（多个文件路径用换行分隔）"
 
-#: main.py:68
+#: main.py:112
 msgid "No test case is selected. Quit generating tests."
 msgstr "未选择测试用例，退出生成单元测试。"
 
-#: main.py:90
-msgid "No valid file is provided. Will not use reference to generate tests."
-msgstr "没有提供合法的文件，生成单元测试时将不使用参考。"
+#: main.py:129
+msgid "Will generate tests for the following cases."
+msgstr "将为以下用例生成测试。"
 
-#: main.py:93
-msgid "Will use the following reference files to generate tests:"
-msgstr "将参考以下文件生成单元测试："
+#: main.py:132
+msgid ""
+"\n"
+"Test cases:"
+msgstr "\n测试用例："
 
-#: main.py:94
+#: main.py:139
+msgid ""
+"\n"
+"No valid reference file is provided. Will not use reference to generate "
+"tests."
+msgstr "\n没有提供合法的文件，生成单元测试时将不使用参考。"
+
+#: main.py:144
+msgid ""
+"\n"
+"Will use the following reference files to generate tests."
+msgstr "\n将参考以下文件生成单元测试。"
+
+#: main.py:145
 msgid ""
 "\n"
 "Valid reference files:"
@@ -52,7 +75,7 @@ msgstr ""
 "\n"
 "参考文件："
 
-#: main.py:98
+#: main.py:150
 msgid ""
 "\n"
 "Invalid files:"
@@ -60,7 +83,7 @@ msgstr ""
 "\n"
 "不合法的文件："
 
-#: main.py:169
+#: main.py:232
 msgid "The function's size surpasses AI's context capacity."
 msgstr "由于当前函数过大，AI暂时无法处理。"
 
diff --git a/unit_tests/main.py b/unit_tests/main.py
index b17a4cc..f55a6c9 100644
--- a/unit_tests/main.py
+++ b/unit_tests/main.py
@@ -1,5 +1,6 @@
 import os
 import sys
+from typing import List, Tuple
 
 import click
 from find_reference_tests import find_reference_tests
@@ -15,89 +16,162 @@ from chatmark import Checkbox, Form, Step, TextEditor  # noqa: E402
 from ide_services import ide_language  # noqa: E402
 
 
-def generate_unit_tests_workflow(
-    user_prompt: str,
-    func_to_test: FuncToTest,
-    tui_lang: TUILanguage,
-):
-    """
-    The main workflow for generating unit tests.
-    """
-    repo_root = os.getcwd()
+class UnitTestsWorkflow:
+    def __init__(
+        self,
+        user_prompt: str,
+        func_to_test: FuncToTest,
+        repo_root: str,
+        tui_lang: TUILanguage,
+    ):
+        self.user_prompt = user_prompt
+        self.func_to_test = func_to_test
+        self.repo_root = repo_root
+        self.tui_lang = tui_lang
 
-    _i = get_translation(tui_lang)
+    def run(self):
+        """
+        Run the workflow to generate unit tests.
+        """
+        cases, files = self.step1_propose_cases_and_reference_files()
 
-    msg = _i("Analyzing the function and current unit tests...")
+        cases, files = self.step2_edit_cases_and_reference_files(cases, files)
 
-    with Step(msg):
-        test_cases = propose_test(
-            user_prompt=user_prompt,
-            func_to_test=func_to_test,
-            chat_language=tui_lang.chat_language,
+        self.step3_write_and_print_tests(cases, files)
+
+    def step1_propose_cases_and_reference_files(self) -> Tuple[List[str], List[str]]:
+        """
+        Propose test cases and reference files for a specified function.
+
+        Return: (test_cases, reference_files)
+        """
+        test_cases: List[str] = []
+        reference_files: List[str] = []
+
+        _i = get_translation(self.tui_lang)
+
+        msg = _i("Analyzing the function and current unit tests...")
+
+        with Step(msg):
+            test_cases = propose_test(
+                user_prompt=self.user_prompt,
+                func_to_test=self.func_to_test,
+                chat_language=self.tui_lang.chat_language,
+            )
+
+            ref_files = find_reference_tests(
+                self.repo_root, self.func_to_test.func_name, self.func_to_test.file_path
+            )
+
+        if ref_files:
+            # Only use the most relevant reference file currently
+            reference_files.append(ref_files[0])
+
+        return test_cases, reference_files
+
+    def step2_edit_cases_and_reference_files(
+        self, test_cases: List[str], reference_files: List[str]
+    ) -> Tuple[List[str], List[str]]:
+        """
+        Edit test cases and reference files by user.
+
+        Return the updated cases and valid reference files.
+        """
+        _i = get_translation(self.tui_lang)
+
+        checkbox = Checkbox(
+            options=test_cases,
+            title=_i("Select test cases to generate"),
+        )
+        case_editor = TextEditor(
+            text="",
+            title=_i(
+                "You can add more test cases here\n"
+                "(Multiple cases can be separated by line breaks)"
+            ),
+        )
+        ref_editor = TextEditor(
+            text=reference_files[0] if reference_files else "",
+            title=_i("Edit reference test file\n(Multiple files can be separated by line breaks)"),
         )
 
-        ref_files = find_reference_tests(repo_root, func_to_test.func_name, func_to_test.file_path)
+        form = Form(components=[checkbox, case_editor, ref_editor])
+        form.render()
 
-    ref_file = ref_files[0] if ref_files else ""
+        # Check test cases
+        cases = [checkbox.options[idx] for idx in checkbox.selections]
+        user_cases = []
+        if case_editor.new_text:
+            user_cases = [c.strip() for c in case_editor.new_text.split("\n")]
+            user_cases = [c for c in user_cases if c]
 
-    cases_checkbox = Checkbox(
-        options=test_cases,
-        title=_i("Select test cases to generate"),
-    )
-    ref_file_editor = TextEditor(
-        text=ref_file,
-        title=_i("Edit reference test file\n(Multiple files can be separated by line breaks)"),
-    )
+        cases.extend(user_cases)
 
-    form = Form(components=[cases_checkbox, ref_file_editor])
-    form.render()
+        # Check if any test case is selected
+        if not cases:
+            raise UserCancelledException(_i("No test case is selected. Quit generating tests."))
 
-    selected_cases = [cases_checkbox.options[idx] for idx in cases_checkbox.selections]
-    new_refs = ref_file_editor.new_text
+        # Validate reference files
+        ref_files = [f.strip() for f in ref_editor.new_text.split("\n")]
+        valid_files = []
+        invalid_files = []
 
-    # Check user input
-    # Check if any test case is selected
-    if not cases_checkbox.selections:
-        raise UserCancelledException(_i("No test case is selected. Quit generating tests."))
+        for ref_file in ref_files:
+            if not ref_file:
+                continue
+            try:
+                retrieve_file_content(file_path=ref_file, root_path=self.repo_root)
+                valid_files.append(ref_file)
+            except Exception as e:
+                invalid_files.append(ref_file)
 
-    # Validate reference files
-    valid_files = []
-    invalid_files = []
-    ref_files = [f.strip() for f in new_refs.split("\n")]
+        # Print summary
+        title = _i("Will generate tests for the following cases.")
+        lines = []
 
-    for ref_file in ref_files:
-        if not ref_file:
-            continue
-        try:
-            retrieve_file_content(file_path=ref_file, root_path=repo_root)
-            valid_files.append(ref_file)
+        lines.append(_i("\nTest cases:"))
+        width = len(str(len(cases)))
+        lines.extend([f"{(i+1):>{width}}. {c}" for i, c in enumerate(cases)])
 
-        except Exception as e:
-            invalid_files.append(ref_file)
+        if not valid_files:
+            lines.append(
+                _i(
+                    "\nNo valid reference file is provided. "
+                    "Will not use reference to generate tests."
+                )
+            )
+        else:
+            lines.append(_i("\nWill use the following reference files to generate tests."))
+            lines.append(_i("\nValid reference files:"))
+            width = len(str(len(valid_files)))
+            lines.extend([f"{(i+1):>{width}}. {f}" for i, f in enumerate(valid_files)])
 
-    title = ""
-    lines = []
-    if not valid_files:
-        title = _i("No valid file is provided. Will not use reference to generate tests.")
-    else:
-        title = _i("Will use the following reference files to generate tests:")
-        lines.append(_i("\nValid reference files:"))
-        lines.extend(valid_files)
+        if invalid_files:
+            lines.append(_i("\nInvalid files:"))
+            width = len(str(len(invalid_files)))
+            lines.extend([f"{(i+1):>{width}}. {f}" for i, f in enumerate(invalid_files)])
 
-    if invalid_files:
-        lines.append(_i("\nInvalid files:"))
-        lines.extend(invalid_files)
+        with Step(title):
+            print("\n".join(lines), flush=True)
 
-    with Step(title):
-        print("\n".join(lines), flush=True)
+        return cases, valid_files
 
-    write_and_print_tests(
-        root_path=repo_root,
-        func_to_test=func_to_test,
-        test_cases=selected_cases,
-        reference_files=valid_files,
-        chat_language=tui_lang.chat_language,
-    )
+    def step3_write_and_print_tests(
+        self,
+        cases: List[str],
+        ref_files: List[str],
+    ):
+        """
+        Write and print tests.
+        """
+
+        write_and_print_tests(
+            root_path=self.repo_root,
+            func_to_test=self.func_to_test,
+            test_cases=cases,
+            reference_files=ref_files,
+            chat_language=self.tui_lang.chat_language,
+        )
 
 
 @click.command()
@@ -150,7 +224,8 @@ def main(input: str):
     )
 
     try:
-        generate_unit_tests_workflow(user_prompt, func_to_test, tui_lang)
+        workflow = UnitTestsWorkflow(user_prompt, func_to_test, repo_root, tui_lang)
+        workflow.run()
 
     except TokenBudgetExceededException as e:
         msg = _i("The function's size surpasses AI's context capacity.")