rentgen/tests/test_verify.py

#!/usr/bin/env python3
"""
test_verify.py - Extension badge count verification test

Verifies Rentgen's core functionality: counting third-party domains.
Tests two scenarios:
1. Site without third-party domains (news.ycombinator.com) → badge = 0
2. Site with third-party domains (pudelek.pl) → badge > 0
"""

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_badge_count():
    """Test badge count for sites with/without third-party domains. Returns (success, results_dict)."""
    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()

        # Load test library
        test_lib_path = os.path.join(os.path.dirname(__file__), 'test-lib.js')
        with open(test_lib_path, 'r') as f:
            test_lib = f.read()

        results = {}

        # Test 1: Site without third-party domains
        print("  Testing news.ycombinator.com (expected: 0 third-party domains)...")
        client.navigate("https://news.ycombinator.com")
        time.sleep(5)  # Wait for page load and tracking detection

        badge_count_hn = client.execute_script(
            test_lib + "\nreturn testBadgeCount();",
            script_timeout=10000
        )
        results['hn'] = badge_count_hn
        print(f"    → Badge count: {badge_count_hn}")

        # Test 2: Site with third-party domains
        print("  Testing pudelek.pl (expected: >0 third-party domains)...")
        client.navigate("https://pudelek.pl")
        time.sleep(10)  # Wait longer for page load and tracking detection

        badge_count_pudelek = client.execute_script(
            test_lib + "\nreturn testBadgeCount();",
            script_timeout=10000
        )
        results['pudelek'] = badge_count_pudelek
        print(f"    → Badge count: {badge_count_pudelek}")

        client.close()

        # Verify results
        if badge_count_hn is None or badge_count_pudelek is None:
            return False, "Test timed out or failed to get badge count"

        if badge_count_hn == 0 and badge_count_pudelek > 0:
            return True, results
        else:
            return False, f"Unexpected results: HN={badge_count_hn} (expected 0), Pudelek={badge_count_pudelek} (expected >0)"

    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."""
    print("Starting Rentgen badge count verification test...")
    print()

    xvfb_pid = start_xvfb()
    webext_pid = start_webext()

    success, result = test_badge_count()

    cleanup(xvfb_pid, webext_pid)

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

    print(f"PASS: Badge count test succeeded!")
    print(f"  - news.ycombinator.com: {result['hn']} third-party domains")
    print(f"  - pudelek.pl: {result['pudelek']} third-party domains")
    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)