Update backend.py

backend.py

@@ -3,7 +3,9 @@ import io
 import requests
 import html  # For escaping HTML characters
 from bs4 import BeautifulSoup
-import pandas as pd  # Added pandas for Excel export
+import pandas as pd
+from openpyxl import Workbook
+from openpyxl.styles import Alignment, Font
 from openai import OpenAI
 
 # Initialize OpenAI API with Nvidia's Llama model
@@ -24,26 +26,19 @@ def clean_test_case_output(text):
 def generate_testcases(user_story):
     """
     Generates advanced QA test cases based on a provided user story by interacting 
-    with Nvidia's llama model API. The prompt is refined for clarity, 
-    and the output is processed for better quality.
+    with Nvidia's llama model API.
     
     :param user_story: A string representing the user story for which to generate test cases.
-    :return: A list of test cases in the form of dictionaries.
+    :return: A list of dictionaries with test case information.
     """
-
-    # Few-shot learning examples to guide the model
     few_shot_examples = """
-    "if its not a DropBury or ODAC Portal User Story, then we perform testing in Tech360 iOS App"
-    "Generate as many as testcases possible minimum 6 ,maximum it can be anything"
-    "Understand the story thoroughly"
-    "If it's a DropBury or ODAC Portal User Story, then we perform testing in ODAC Portal"
+    "Generate as many test cases as possible. Minimum 6, but can be more."
+    "Structure each test case with Test Case, Preconditions, Steps, and Expected Result."
     """
 
-    # Combine the few-shot examples with the user story for the model to process
     prompt = few_shot_examples + f"\nUser Story: {user_story}\n"
 
     try:
-        # Call the Nvidia llama API with the refined prompt
         completion = client.chat.completions.create(
             model="meta/llama-3.1-405b-instruct",
             messages=[
@@ -55,37 +50,65 @@ def generate_testcases(user_story):
             stream=True
         )
 
-        # Initialize an empty string to accumulate the response
         test_cases_text = ""
 
-        # Accumulate the response from the streaming chunks
         for chunk in completion:
             if chunk.choices[0].delta.content is not None:
                 test_cases_text += chunk.choices[0].delta.content
 
-        # Ensure the entire response is captured before cleaning
         if test_cases_text.strip() == "":
             return [{"test_case": "No test cases generated or output was empty."}]
 
-        # Clean the output by unescaping HTML entities and replacing <br> tags
         test_cases_text = clean_test_case_output(test_cases_text)
 
         try:
-            # Try to parse the output as JSON, assuming the model returns structured test cases
             test_cases = json.loads(test_cases_text)
             if isinstance(test_cases, list):
-                return test_cases  # Return structured test cases
+                return test_cases
             else:
-                return [{"test_case": test_cases_text}]  # Return as a list with the text wrapped in a dict
-
+                return [{"test_case": test_cases_text}]
         except json.JSONDecodeError:
-            # Fallback: return the raw text if JSON parsing fails
-            return [{"test_case": test_cases_text}]
+            # If JSON decoding fails, attempt to parse manually
+            return parse_test_cases(test_cases_text)
 
     except requests.exceptions.RequestException as e:
         print(f"API request failed: {str(e)}")
         return []
 
+def parse_test_cases(raw_text):
+    """
+    Parse raw text output into structured test cases.
+    
+    :param raw_text: Raw text returned from the model.
+    :return: List of dictionaries with structured test cases.
+    """
+    test_cases = []
+    case = {
+        "Test Case": "",
+        "Preconditions": "N/A",
+        "Steps": "N/A",
+        "Expected Result": "N/A"
+    }
+
+    lines = raw_text.split("\n")
+    for line in lines:
+        if line.startswith("Test Case"):
+            if case["Test Case"]:
+                test_cases.append(case)  # Save the previous test case
+                case = {"Test Case": "", "Preconditions": "N/A", "Steps": "N/A", "Expected Result": "N/A"}
+            case["Test Case"] = line.replace("Test Case:", "").strip()
+        elif "Preconditions" in line:
+            case["Preconditions"] = line.replace("Preconditions:", "").strip() or "N/A"
+        elif "Steps" in line:
+            case["Steps"] = line.replace("Steps:", "").strip() or "N/A"
+        elif "Expected Result" in line:
+            case["Expected Result"] = line.replace("Expected Result:", "").strip() or "N/A"
+    
+    if case["Test Case"]:  # Add the last case
+        test_cases.append(case)
+
+    return test_cases
+
 def export_test_cases(test_cases):
     """
     Exports the test cases to an Excel file with specific columns:
@@ -94,39 +117,21 @@ def export_test_cases(test_cases):
     - Steps
     - Expected Result
     
-    :param test_cases: A list of test case dictionaries or raw text.
+    :param test_cases: A list of test case dictionaries.
     :return: Bytes of the Excel file.
     """
     if not test_cases:
         return "No test cases to export."
 
-    # Define the structure of the Excel file
     formatted_test_cases = []
     
     for case in test_cases:
-        # Assuming each test case is a dictionary with 'test_case' content or similar
-        test_case_content = case.get('test_case', '')
-        
-        # Split the content into separate sections (you might need to adjust based on actual output structure)
-        lines = test_case_content.split('\n')
-        test_case = ""
-        preconditions = ""
-        steps = ""
-        expected_result = ""
-        
-        for line in lines:
-            if "Preconditions" in line:
-                preconditions = line.replace("Preconditions:", "").strip()
-            elif "Steps" in line:
-                steps = line.replace("Steps:", "").strip()
-            elif "Expected Result" in line:
-                expected_result = line.replace("Expected Result:", "").strip()
-            else:
-                # Default to putting the first part as the "Test Case"
-                if not test_case:
-                    test_case = line.strip()
-        
-        # Append to formatted test cases list
+        # Ensure each field has a default value if missing
+        test_case = case.get('Test Case', 'N/A')
+        preconditions = case.get('Preconditions', 'N/A')
+        steps = case.get('Steps', 'N/A')
+        expected_result = case.get('Expected Result', 'N/A')
+
         formatted_test_cases.append({
             'Test Case': test_case,
             'Preconditions': preconditions,
@@ -134,12 +139,30 @@ def export_test_cases(test_cases):
             'Expected Result': expected_result
         })
     
-    # Convert the list of dictionaries into a DataFrame
-    df = pd.DataFrame(formatted_test_cases)
+    wb = Workbook()
+    ws = wb.active
+    ws.title = "Test Cases"
+
+    # Add headers with formatting
+    headers = ["Test Case", "Preconditions", "Steps", "Expected Result"]
+    ws.append(headers)
     
-    # Create an Excel file using pandas
-    output = io.BytesIO()
-    df.to_excel(output, index=False, engine='openpyxl')  # Export to Excel without index
-    output.seek(0)  # Rewind the buffer to the beginning
+    for cell in ws[1]:
+        cell.font = Font(bold=True)
+        cell.alignment = Alignment(horizontal="center", vertical="center")
+
+    # Add the test case data
+    for case in formatted_test_cases:
+        ws.append([case["Test Case"], case["Preconditions"], case["Steps"], case["Expected Result"]])
     
+    # Adjust column widths for neatness
+    ws.column_dimensions['A'].width = 50  # Test Case
+    ws.column_dimensions['B'].width = 30  # Preconditions
+    ws.column_dimensions['C'].width = 50  # Steps
+    ws.column_dimensions['D'].width = 50  # Expected Result
+
+    output = io.BytesIO()
+    wb.save(output)
+    output.seek(0)
+
     return output.getvalue()