"""Regression checks for Playwright ownership across Gradio/ZeroGPU workers.

Run with: python -m unittest discover -s tests
Requires the Space's browser dependencies and installed Chromium, but no models.
"""
import multiprocessing
import unittest
from concurrent.futures import ThreadPoolExecutor

import browsergym.core

from browser_session import browser_session


class StaleConnection:
    def __getattr__(self, name):
        raise AssertionError('Reused a connection from another event loop')


def exercise_browser():
    for _ in range(2):
        with browser_session():
            playwright = browsergym.core._get_global_playwright()
            playwright.selectors.set_test_id_attribute('bid')
            browser = playwright.chromium.launch(headless=True)
            try:
                page = browser.new_page()
                page.set_content('<button bid="13">Click Me!</button>')
                page.get_by_test_id('13').click()
            finally:
                browser.close()
        assert browsergym.core._PLAYWRIGHT is None


class BrowserSessionTests(unittest.TestCase):
    def tearDown(self):
        browsergym.core._set_global_playwright(None)

    def test_reused_thread(self):
        browsergym.core._set_global_playwright(StaleConnection())
        with ThreadPoolExecutor(max_workers=1) as pool:
            pool.submit(exercise_browser).result(timeout=30)
            pool.submit(exercise_browser).result(timeout=30)

    def test_forked_worker(self):
        browsergym.core._set_global_playwright(StaleConnection())
        process = multiprocessing.get_context('fork').Process(target=exercise_browser)
        process.start()
        process.join(timeout=30)
        if process.is_alive():
            process.terminate()
            process.join()
            self.fail('Browser worker timed out')
        self.assertEqual(process.exitcode, 0)

    def test_error_cleanup(self):
        with self.assertRaisesRegex(ValueError, 'example'):
            with browser_session():
                raise ValueError('example')
        self.assertIsNone(browsergym.core._PLAYWRIGHT)