Update app.py

app.py

@@ -272,9 +272,17 @@ class TestCaseGenerator:
                     self.functions[method_name] = ''.join(method_body)
                     method_body = []
 
+    # def generate_response(self, prompt):
+    #     response = [
+    #         ChatMessage(role="system", content="You are a sincere and helpful coding assistant"),
+    #         ChatMessage(role="user", content=prompt),
+    #     ]
+    #     resp = MistralAI(self.model).chat(response)
+    #     return resp
+
     def generate_response(self, prompt):
         response = [
-            ChatMessage(role="system", content="You are a sincere and helpful coding assistant"),
+            ChatMessage(role="system", content="You are tasked with generating unit test cases, including descriptions and assert statements, for a given function. You will also calculate the test coverage percentage. Follow these instructions carefully:"),
             ChatMessage(role="user", content=prompt),
         ]
         resp = MistralAI(self.model).chat(response)
@@ -286,26 +294,64 @@ class TestCaseGenerator:
             return
 
         snippet = self.functions[func]
-        prompt = f"""Your task is to generate unit test cases with maximum coverage for this function : \n\n{snippet}\
-        \n\n - Generate unique unit test cases. Include couple of edge cases as well.
-        \n\n - Calculate the coverage percentage and display in the response.
-        \n\n - All the test cases should have the mandatory assert statement.
-        \n\n - Every test case should be defined as a method inside the class.
-        \n\n - All the test cases should have textual description.
-        \n\n - Politely refuse if the function is not suitable for generating test cases.
-        \n\n - There should be no duplicate and incomplete test case.
-        \n\n - Avoid generating repeated statements.
-        \n\n - Recheck your response before generating.
-        \n\n - Do not share the last Test Case.
+        , file_extension = os.path.splitext(uploaded_file.name)
+        lang = file_extension[1:]
+       
+        prompt = f"""1. You will be provided with the following inputs:
+                           <function_code>
+                           {snippet}
+                           </function_code>
+
+                           <language>{lang}</language> 
+
+                    2. Analyze the provided function:
+                       - Identify the function name, parameters, and return type
+                       - Determine the main logic and branches in the function
+                       - Note any potential edge cases or boundary conditions
+
+                    3. Generate unit test cases:
+                       - Create at least 3-5 test cases that cover different scenarios
+                       - Include normal cases, edge cases, and potential error conditions
+                       - Ensure that the test cases collectively cover all branches and logic paths in the function
+
+                    4. For each test case, provide:
+                       - A brief description of the test scenario
+                       - Input values for the function parameters
+                       - The expected output or behavior
+                       - An assert statement in the appropriate syntax for the given programming language
+
+                    5. Calculate the test coverage percentage:
+                       - Determine the number of code paths or branches in the function
+                       - Count how many of these paths are covered by your test cases
+                       - Calculate the percentage: (covered paths / total paths) * 100
+
+                    6. Present your output in the following format:
+                           <unit_tests>
+                           <test_case>
+                           <description>Description of the test case</description>
+                           <input>Input values for the function</input>
+                           <expected_output>Expected output or behavior</expected_output>
+                           <assert_statement>Assert statement in the appropriate language syntax</assert_statement>
+                           </test_case>
+                           <!-- Repeat for each test case -->
+                           </unit_tests>
+
+                           <test_coverage>
+                           <percentage>Calculated test coverage percentage</percentage>
+                           <explanation>Brief explanation of how the percentage was calculated</explanation>
+                           </test_coverage>
+
+                    Remember to adapt your assert statements and syntax to the specific programming language provided. If you're unsure about the exact syntax for a particular language, use a general pseudocode format that clearly conveys the assertion logic.
         """
         resp = self.generate_response(prompt)
-        post_prompt = f"""Except the last test case, display everything that is present in this end to end: \n\n{resp}\
-        \n\n - Do not add anything extra. Just copy and paste everything except the last test case.
-        \n\n - Do not mention the count of total number of test cases in the response.
-        \n\n - Do not mention this sentence - "I have excluded the last test case as per your request"
-        """
-        post_resp = self.generate_response(post_prompt)
-        return post_resp
+        
+        # post_prompt = f"""Except the last test case, display everything that is present in this end to end: \n\n{resp}\
+        # \n\n - Do not add anything extra. Just copy and paste everything except the last test case.
+        # \n\n - Do not mention the count of total number of test cases in the response.
+        # \n\n - Do not mention this sentence - "I have excluded the last test case as per your request"
+        # """
+        # post_resp = self.generate_response(post_prompt)
+        return resp
 
     def run(self):
         self.setup_streamlit_ui()