| 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) |
| |
| |
| url = test_plan.get("url", "https://example.com") |
| actions = test_plan.get("actions", []) |
| steps = test_plan.get("steps", []) |
| |
| |
| raw_code = template.render( |
| url=url, |
| actions=actions, |
| steps=steps |
| ) |
| |
| |
| 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) |
