Selenium: Mini Framework

import unittest
import time
import os
import subprocess
import threading
import http.server
import socketserver
import sys
from datetime import datetime
from selenium_base_test import BaseSQLMappingAppTest

# --- Configuration ---
SOURCE_CODE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), 'SourceCode'))
DIST_DIR = os.path.join(SOURCE_CODE_DIR, 'dist')
SERVER_PORT = 8000
SELENIUM_TEST_DIR = os.path.join(os.path.dirname(__file__), "selenium")
MAIN_SCREENSHOT_DIR = os.path.join(os.path.dirname(__file__), "temp/selenium", datetime.now().strftime("%Y%m%d_%H%M%S"))

# --- Global Variables ---
httpd_server = None

def start_server():
    global httpd_server
    class Handler(http.server.SimpleHTTPRequestHandler):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, directory=DIST_DIR, **kwargs)
    httpd_server = socketserver.TCPServer(("", SERVER_PORT), Handler)
    print(f"\nStarting local server for '{DIST_DIR}' at http://localhost:{SERVER_PORT}")
    httpd_server.serve_forever()

class ScreenshotTestResult(unittest.TextTestResult):
    def __init__(self, stream, descriptions, verbosity, screenshot_dir=None):
        super().__init__(stream, descriptions, verbosity)
        self.screenshot_dir = screenshot_dir

    def addError(self, test, err):
        super().addError(test, err)
        # Set a flag on the test instance to prevent success screenshot
        setattr(test, '_test_has_failed', True)
        driver = getattr(test, 'driver', None)
        
        # This block now runs ALWAYS on error, regardless of --log-all
        if driver and self.screenshot_dir:
            test_method_name = test.id().split('.')[-1]
            print(f"\n--- Test '{test_method_name}' FAILED. Performing error handling. ---")
            
            screenshot_path = os.path.join(self.screenshot_dir, f"{test_method_name}_ERROR.png")
            driver.save_screenshot(screenshot_path)
            print(f"--- Saved error screenshot to: {screenshot_path} ---")

            # FIX: Change browser log filename on error
            browser_log_path = os.path.join(self.screenshot_dir, f"{test_method_name}_browser_ERROR.log")
            try:
                browser_logs = driver.get_log('browser')
                with open(browser_log_path, 'w') as f:
                    f.write(f"==== Browser Console Logs for {test_method_name} ====\n")
                    if browser_logs:
                        for entry in browser_logs: f.write(f"[{entry['level']}] {entry['message']}\n")
                    else:
                        f.write("No browser console logs were found.\n")
                print(f"--- Browser console logs saved to: {browser_log_path} ---")
            except Exception as e:
                with open(browser_log_path, 'w') as f: f.write(f"Could not retrieve browser logs: {e}\n")
                print(f"--- Error retrieving browser logs. See: {browser_log_path} ---")
            
            print(f"--- BROWSER FOR '{test_method_name}' WILL REMAIN OPEN for inspection. ---")
            setattr(test, 'keep_browser_open', True)

class CustomTestRunner(unittest.TextTestRunner):
    def __init__(self, *args, **kwargs):
        self.screenshot_dir = kwargs.pop('screenshot_dir', None)
        super().__init__(*args, **kwargs)

    def _makeResult(self):
        return ScreenshotTestResult(self.stream, self.descriptions, self.verbosity, self.screenshot_dir)

def run_tests():
    server_thread = threading.Thread(target=start_server)
    server_thread.daemon = True
    server_thread.start()
    time.sleep(1)

    try:
        BaseSQLMappingAppTest.SCREENSHOT_DIR = MAIN_SCREENSHOT_DIR
        if '--log-all' in sys.argv:
            BaseSQLMappingAppTest.LOG_ALL = True
            sys.argv.remove('--log-all')
            print("--- Log-all mode enabled. A screenshot will be taken for every test. ---")

        print(f"\n--- Discovering tests in: {SELENIUM_TEST_DIR} ---")
        loader = unittest.TestLoader()
        suite = loader.discover(start_dir=SELENIUM_TEST_DIR, pattern="test_*.py")
        
        log_file_path = os.path.join(MAIN_SCREENSHOT_DIR, 'test.log')
        print(f"--- Starting test run. Full log will be saved to: {log_file_path} ---")
        with open(log_file_path, 'w') as log_file:
            runner = CustomTestRunner(
                stream=log_file, 
                verbosity=2,
                screenshot_dir=MAIN_SCREENSHOT_DIR
            )
            result = runner.run(suite)
        
        print(f"--- Test run complete. Full log saved to: {log_file_path} ---")
        print(f"--- All artifacts are in: {os.path.abspath(MAIN_SCREENSHOT_DIR)} ---")
        
        if not result.wasSuccessful():
            sys.exit(1)

    finally:
        global httpd_server
        if httpd_server:
            print("\n--- Shutting down local server ---")
            httpd_server.shutdown()
            httpd_server.server_close()

if __name__ == "__main__":
    os.makedirs(MAIN_SCREENSHOT_DIR, exist_ok=True)
    run_tests()

import unittest
import time
import os
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait, Select
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.action_chains import ActionChains
from selenium.common.exceptions import TimeoutException

# --- Configuration (can be accessed by all tests) ---
APP_URL = "http://localhost:8000"

class BaseSQLMappingAppTest(unittest.TestCase):
    """
    This is the base class for all test cases.
    It contains the common setup, teardown, and helper methods.
    """
    LOG_ALL = False
    SCREENSHOT_DIR = "" # This will be set by the test runner
    _test_has_failed = False # Internal flag to track test failure

    def setUp(self):
        self._test_has_failed = False
        self.keep_browser_open = False
        chrome_options = Options()
        chrome_options.set_capability('goog:loggingPrefs', {'browser': 'ALL'})
        self.driver = webdriver.Chrome(options=chrome_options)
        self.driver.get(APP_URL)
        # FIX: Increased timeout to 30 seconds
        self.wait = WebDriverWait(self.driver, 30)

    def tearDown(self):
        # FIX: This is the new, correct logic for --log-all.
        # If log-all is enabled AND the test has NOT failed, take a success screenshot.
        # This block will be SKIPPED if the test failed, preserving the _ERROR.png file.
        if self.LOG_ALL and not self._test_has_failed:
            test_method_name = self.id().split('.')[-1]
            screenshot_path = os.path.join(self.SCREENSHOT_DIR, f"{test_method_name}_SUCCESS.png")
            self.driver.save_screenshot(screenshot_path)
            print(f"--- Saved success screenshot for '{test_method_name}' to: {screenshot_path} ---")

        if not self.keep_browser_open:
            self.driver.quit()

    def wait_for_editor(self, editor_name):
        try:
            container_id = "SQLSourceIDE" if editor_name == "sqlSourceEditor" else "SQLOutputIDE"
            self.wait.until(EC.presence_of_element_located((By.ID, container_id)))
            self.wait.until(
                lambda driver: driver.execute_script(
                    f"return window.monacoEditors && window.monacoEditors.{editor_name}"
                )
            )
        except TimeoutException:
            print(f"Timeout waiting for Monaco Editor instance '{editor_name}' to be created.")
            raise

    def get_monaco_editor_text(self, editor_name):
        self.wait_for_editor(editor_name)
        js_script = f"return window.monacoEditors.{editor_name}.getValue();"
        return self.driver.execute_script(js_script)

    def set_monaco_editor_text(self, editor_name, text):
        self.wait_for_editor(editor_name)
        escaped_text = text.replace('\\', '\\\\').replace("'", "\\'").replace('"', '\\"').replace('\n', '\\n').replace('`', '\\`')
        js_script = f"window.monacoEditors.{editor_name}.setValue(`{escaped_text}`);"
        self.driver.execute_script(js_script)

from selenium_base_test import BaseSQLMappingAppTest

class TestInitialLoad(BaseSQLMappingAppTest):
    def test_01_initial_load_and_default_sql(self):
        print("\nRunning test_01_initial_load_and_default_sql...")
        editor_text = self.get_monaco_editor_text("sqlSourceEditor")
        self.assertIn("INSERT INTO Customers", editor_text)
        print("OK")