backend.py

import csv
import io
from transformers import pipeline
from bs4 import BeautifulSoup

# Initialize the model (e.g., FLAN-T5)
model = pipeline(
    "text-generation",
    model="google/flan-t5-xxl",  # Replace with your chosen model
    torch_dtype=torch.float16  # Reduce memory usage
)

def clean_test_case_output(text):
    text = text.strip()
    soup = BeautifulSoup(text, 'html.parser')
    return soup.get_text(separator="\n").strip()

def generate_testcases(user_story):
    few_shot_examples = """
    "Generate as many test cases as possible (minimum 6)"
    "Understand the user story thoroughly"
    "If it's a portal-related user story, test in ODAC Portal; else, test in Tech360 iOS App"
    """
    
    prompt = f"{few_shot_examples}\nUser Story: {user_story}"
    
    try:
        result = model(prompt, max_length=4096)[0]["generated_text"]
        cleaned_text = clean_test_case_output(result)
        
        try:
            test_cases = json.loads(cleaned_text)
            if isinstance(test_cases, list):
                return test_cases
            else:
                return [{"test_case": cleaned_text}]
        except json.JSONDecodeError:
            return [{"test_case": cleaned_text}]
    
    except Exception as e:
        return [{"test_case": f"Error: {str(e)}"}]

def export_test_cases(test_cases, format='csv'):
    if not test_cases:
        return "No test cases to export."
    
    if format == 'csv':
        output = io.StringIO()
        writer = csv.writer(output)
        writer.writerow(["Test Case"])
        for case in test_cases:
            writer.writerow([case.get("test_case", "N/A")])
        return output.getvalue()
    else:
        raise ValueError("Unsupported format. Only CSV is supported.")

def save_test_cases_as_file(test_cases, format='csv'):
    if format == 'csv':
        with open('test_cases.csv', 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(["Test Case"])
            for case in test_cases:
                writer.writerow([case.get("test_case", "N/A")])
        return "test_cases.csv saved"
    else:
        return f"Unsupported format: {format}"