sd-webui-deforum/tests/conftest.py

59 lines
2.4 KiB
Python

import pytest
import subprocess
import sys
import os
from subprocess import Popen, PIPE, STDOUT
from pathlib import Path
from tenacity import retry, stop_after_delay, wait_fixed
import threading
import requests
def pytest_addoption(parser):
parser.addoption("--start-server", action="store_true", help="start the server before the test run (if not specified, you must start the server manually)")
@pytest.fixture
def cmdopt(request):
return request.config.getoption("--start-server")
@retry(wait=wait_fixed(5), stop=stop_after_delay(60))
def wait_for_service(url):
response = requests.get(url, timeout=(5, 5))
print(f"Waiting for server to respond 200 at {url} (response: {response.status_code})...")
assert response.status_code == 200
@pytest.fixture(scope="session", autouse=True)
def start_server(request):
if request.config.getoption("--start-server"):
# Kick off server subprocess
script_directory = os.path.dirname(__file__)
a1111_directory = Path(script_directory).parent.parent.parent # sd-webui/extensions/deforum/tests/ -> sd-webui
print(f"Starting server in {a1111_directory}...")
proc = Popen(["python", "-m", "coverage", "run", "--data-file=.coverage.server", "launch.py",
"--skip-prepare-environment", "--skip-torch-cuda-test", "--test-server", "--no-half",
"--disable-opt-split-attention", "--use-cpu", "all", "--add-stop-route", "--api", "--deforum-api", "--listen"],
cwd=a1111_directory,
stdout=PIPE,
stderr=STDOUT,
universal_newlines=True)
# ensure server is killed at the end of the test run
request.addfinalizer(proc.kill)
# Spin up separate thread to capture the server output to file and stdout
def server_console_manager():
with proc.stdout, open('serverlog.txt', 'ab') as logfile:
for line in proc.stdout:
sys.stdout.write(f"[SERVER LOG] {line}")
sys.stdout.flush()
logfile.write(line.encode('utf-8'))
logfile.flush()
proc.wait()
threading.Thread(target=server_console_manager).start()
# Wait for deforum API to respond
wait_for_service('http://localhost:7860/deforum_api/jobs/')
else:
print("Assuming server is already running...")