File size: 2,978 Bytes
f96a7ac b1a0a54 f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e f96a7ac 8e7043e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | 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) |