nicegui/tests/screen.py

import threading
import time

import pytest
from bs4 import BeautifulSoup
from nicegui import globals, ui
from selenium import webdriver
from selenium.common.exceptions import NoSuchElementException
from selenium.webdriver.remote.webelement import WebElement

PORT = 3392
IGNORED_CLASSES = ['row', 'column', 'q-card', 'q-field', 'q-field__label', 'q-input']


class Screen():

    def __init__(self, selenium: webdriver.Chrome) -> None:
        self.selenium = selenium
        self.server_thread = None

    def start_server(self) -> None:
        '''Start the webserver in a separate thread. This is the equivalent of `ui.run()` in a normal script.'''
        self.server_thread = threading.Thread(target=ui.run, kwargs={'port': PORT, 'show': False, 'reload': False})
        self.server_thread.start()

    def stop_server(self) -> None:
        '''Stop the webserver.'''
        self.selenium.close()
        globals.server.should_exit = True
        self.server_thread.join()

    def open(self, path: str) -> None:
        if self.server_thread is None:
            self.start_server()
        start = time.time()
        while True:
            try:
                self.selenium.get(f'http://localhost:{PORT}{path}')
                break
            except Exception:
                if time.time() - start > 3:
                    raise
                time.sleep(0.1)
                if not self.server_thread.is_alive():
                    raise RuntimeError('The NiceGUI server has stopped running')

    def should_contain(self, text: str) -> None:
        assert self.selenium.title == text or self.find(text), \
            f'could not find "{text}" on:\n{self.render_content()}'

    def should_not_contain(self, text: str) -> None:
        assert self.selenium.title != text
        with pytest.raises(AssertionError):
            element = self.find(text)
            print(element.get_attribute('outerHTML'))

    def click(self, target_text: str) -> None:
        self.find(target_text).click()

    def find(self, text: str) -> WebElement:
        try:
            return self.selenium.find_element_by_xpath(f'//*[contains(text(),"{text}")]')
        except NoSuchElementException:
            raise AssertionError(f'Could not find "{text}" on:\n{self.render_content()}')

    def render_content(self, with_extras: bool = False) -> str:
        body = self.selenium.find_element_by_tag_name('body').get_attribute('innerHTML')
        soup = BeautifulSoup(body, 'html.parser')
        self.simplify_input_tags(soup)
        content = ''
        for child in soup.find_all():
            is_element = False
            if child is None or child.name == 'script':
                continue
            depth = (len(list(child.parents)) - 3) * '  '
            if not child.find_all() and child.text:
                content += depth + child.getText()
                is_element = True
            classes = child.get('class', '')
            if classes:
                if classes[0] in ['row', 'column', 'q-card']:
                    content += depth + classes[0].removeprefix('q-')
                    is_element = True
                if classes[0] == 'q-field':
                    pass
                [classes.remove(c) for c in IGNORED_CLASSES if c in classes]
                for i, c in enumerate(classes):
                    classes[i] = c.removeprefix('q-field--')
                if is_element and with_extras:
                    content += f' [class: {" ".join(classes)}]'

            if is_element:
                content += '\n'

        return f'Title: {self.selenium.title}\n\n{content}'

    @staticmethod
    def simplify_input_tags(soup: BeautifulSoup) -> None:
        for element in soup.find_all(class_="q-field"):
            print(element.prettify())
            new = soup.new_tag('simple_input')
            name = element.find(class_='q-field__label').text
            placeholder = element.find(class_='q-field__native').get('placeholder')
            value = element.find(class_='q-field__native').get('value')
            new.string = (f'{name}: ' if name else '') + (value or placeholder or '')
            new['class'] = element['class']
            element.replace_with(new)

    def get_tags(self, name: str) -> list[WebElement]:
        return self.selenium.find_elements_by_tag_name(name)

    def get_attributes(self, tag: str, attribute: str) -> list[str]:
        return [t.get_attribute(attribute) for t in self.get_tags(tag)]

    def wait(self, t: float) -> None:
        time.sleep(t)