src/core/code_generator.py

from jinja2 import Template
from typing import Dict, Any
from loguru import logger
import black 

class SeleniumGenerator:
    """
    Generates production-grade Selenium scripts.
    Handles missing data gracefully by falling back to comments.
    """

    TEST_TEMPLATE = """
import pytest
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

class GeneratedPage:
    def __init__(self, driver):
        self.driver = driver
        self.url = "{{ url }}"

    def load(self):
        self.driver.get(self.url)

    {% for action in actions %}
    def {{ action.func_name }}(self, value=None):
        # {{ action.description }}
        element = WebDriverWait(self.driver, 10).until(
            EC.presence_of_element_located((By.{{ action.selector_type }}, "{{ action.selector }}"))
        )
        {% if action.type == 'click' %}
        element.click()
        {% elif action.type == 'input' %}
        element.clear()
        element.send_keys(value)
        {% endif %}
    {% endfor %}

@pytest.fixture
def driver():
    options = webdriver.ChromeOptions()
    options.add_argument("--headless") 
    driver = webdriver.Chrome(options=options)
    yield driver
    driver.quit()

def test_generated_scenario(driver):
    page = GeneratedPage(driver)
    page.load()
    
    {% if actions %}
    # Executing structured actions
    {% for action in actions %}
    page.{{ action.func_name }}({{ "'" + action.value + "'" if action.value else '' }})
    {% endfor %}
    {% else %}
    # ⚠️ No selectors provided by AI. Listing manual steps:
    {% for step in steps %}
    # Step: {{ step }}
    {% endfor %}
    {% endif %}
    
    assert driver.title != "" 
"""

    def generate_test_script(self, test_plan: Dict[str, Any]) -> str:
        logger.info("Generating Selenium Script...")
        
        try:
            template = Template(self.TEST_TEMPLATE)
            
            # Safe Data Extraction
            url = test_plan.get("url", "https://example.com")
            actions = test_plan.get("actions", [])
            steps = test_plan.get("steps", [])
            
            # Render code
            raw_code = template.render(
                url=url,
                actions=actions,
                steps=steps
            )
            
            # Try Formatting, but don't crash if it fails
            try:
                return black.format_str(raw_code, mode=black.Mode())
            except Exception as e:
                logger.warning(f"Black formatting failed (returning raw code): {e}")
                return raw_code
            
        except Exception as e:
            logger.error(f"Code Generation Failed: {e}")
            raise e

    def save_script(self, code: str, filename: str = "generated_test.py"):
        with open(filename, "w") as f:
            f.write(code)