backend.py

import json
import io
import requests
import html  # For escaping HTML characters
from bs4 import BeautifulSoup
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
client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key="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"
)

def clean_test_case_output(text):
    """
    Cleans the output to handle HTML characters and unwanted tags.
    """
    text = html.unescape(text)  # Unescape HTML entities
    soup = BeautifulSoup(text, 'html.parser')  # Use BeautifulSoup to handle HTML tags
    cleaned_text = soup.get_text(separator="\n").strip()  # Remove tags and handle newlines
    return cleaned_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.
    
    :param user_story: A string representing the user story for which to generate test cases.
    :return: A list of dictionaries with test case information.
    """
    few_shot_examples = """
    "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."
    """

    prompt = few_shot_examples + f"\nUser Story: {user_story}\n"

    try:
        completion = client.chat.completions.create(
            model="meta/llama-3.1-405b-instruct",
            messages=[
                {"role": "user", "content": prompt}
            ],
            temperature=0.03,
            top_p=0.7,
            max_tokens=4096,
            stream=True
        )

        test_cases_text = ""

        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                test_cases_text += chunk.choices[0].delta.content

        if test_cases_text.strip() == "":
            return [{"test_case": "No test cases generated or output was empty."}]

        test_cases_text = clean_test_case_output(test_cases_text)

        try:
            test_cases = json.loads(test_cases_text)
            if isinstance(test_cases, list):
                return test_cases
            else:
                return [{"test_case": test_cases_text}]
        except json.JSONDecodeError:
            # 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:
    - Test Case
    - Preconditions
    - Steps
    - Expected Result
    
    :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."

    formatted_test_cases = []
    
    for case in test_cases:
        # 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,
            'Steps': steps,
            'Expected Result': expected_result
        })
    
    wb = Workbook()
    ws = wb.active
    ws.title = "Test Cases"

    # Add headers with formatting
    headers = ["Test Case", "Preconditions", "Steps", "Expected Result"]
    ws.append(headers)
    
    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()