rentgen/tests/test_verify.py

#!/usr/bin/env python3
"""
test_verify.py - Minimal extension verification test

Verifies the extension background script is executing by testing
bidirectional communication with a simple addition operation.
"""

import sys
import time
import subprocess
import os
import signal


def is_tty():
    """Check if stdout is a TTY."""
    return sys.stdout.isatty()


def red(text):
    """Return red text if TTY, otherwise plain text."""
    if is_tty():
        return f"\033[91m{text}\033[0m"
    return text


def start_xvfb():
    """Start Xvfb virtual X server. Returns PID."""
    xvfb = subprocess.Popen(
        ["Xvfb", ":99", "-screen", "0", "1024x768x24"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )
    os.environ["DISPLAY"] = ":99"
    time.sleep(2)
    return xvfb.pid


def start_webext():
    """Start web-ext with Marionette enabled. Returns PID."""
    webext = subprocess.Popen(
        ["npx", "web-ext", "run",
         "--arg=-marionette",
         "--arg=--marionette-port",
         "--arg=2828"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    )
    return webext.pid


def test_addition():
    """Test background script via Marionette. Returns (success, result)."""
    try:
        from marionette_driver.marionette import Marionette

        # Wait for Firefox to start
        time.sleep(10)

        # Connect to Marionette
        client = Marionette(host='localhost', port=2828)
        client.start_session()

        # Navigate to any page (needed for content script injection)
        client.navigate("https://example.com")
        time.sleep(5)

        # Test: background should compute (17 * 2) + 3 = 37
        test_value = 17
        expected = 37

        result = client.execute_script("""
            const testValue = arguments[0];

            // Check if content script loaded
            if (!document.body.getAttribute('data-rentgen-injected')) {
                return -1;  // Content script not loaded
            }

            // Dispatch test request to content script
            document.dispatchEvent(new CustomEvent('rentgen_test_request', {
                detail: { value: testValue, timestamp: Date.now() }
            }));

            // Wait for background response
            return new Promise((resolve) => {
                let attempts = 0;
                const checkInterval = setInterval(() => {
                    attempts++;
                    const computed = document.body.getAttribute('data-rentgen-computed');

                    if (computed) {
                        clearInterval(checkInterval);
                        resolve(parseInt(computed));
                    } else if (attempts > 50) {
                        clearInterval(checkInterval);
                        resolve(null);
                    }
                }, 100);
            });
        """, script_args=[test_value], script_timeout=10000)

        client.close()

        if result == expected:
            return True, expected
        else:
            return False, result

    except Exception as e:
        return False, str(e)


def cleanup(xvfb_pid, webext_pid):
    """Kill processes."""
    try:
        os.kill(webext_pid, signal.SIGTERM)
    except:
        pass
    try:
        os.kill(xvfb_pid, signal.SIGTERM)
    except:
        pass


def main():
    """Main test."""
    xvfb_pid = start_xvfb()
    webext_pid = start_webext()

    success, result = test_addition()

    cleanup(xvfb_pid, webext_pid)

    if not success:
        print(red(f"FAIL: Expected 37, got {result}"))
        return 1

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(130)
    except Exception as e:
        print(red(f"ERROR: {e}"))
        sys.exit(1)