mirror of
https://github.com/Abdulazizzn/n8n-enterprise-unlocked.git
synced 2025-12-16 17:46:45 +00:00
100 lines
3.1 KiB
Python
100 lines
3.1 KiB
Python
import asyncio
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from src.constants import (
|
|
ENV_GRACEFUL_SHUTDOWN_TIMEOUT,
|
|
ENV_GRANT_TOKEN,
|
|
ENV_HEALTH_CHECK_SERVER_ENABLED,
|
|
ENV_HEALTH_CHECK_SERVER_PORT,
|
|
ENV_TASK_BROKER_URI,
|
|
ENV_TASK_TIMEOUT,
|
|
)
|
|
|
|
from tests.fixtures.test_constants import (
|
|
GRACEFUL_SHUTDOWN_TIMEOUT,
|
|
HEALTH_CHECK_PORT,
|
|
LOCAL_TASK_BROKER_URL,
|
|
TASK_TIMEOUT,
|
|
)
|
|
|
|
|
|
class TaskRunnerManager:
|
|
"""Responsible for managing the lifecycle of a task runner subprocess."""
|
|
|
|
def __init__(
|
|
self,
|
|
task_broker_url: str = LOCAL_TASK_BROKER_URL,
|
|
graceful_shutdown_timeout: float | None = None,
|
|
):
|
|
self.task_broker_url = task_broker_url
|
|
self.graceful_shutdown_timeout = graceful_shutdown_timeout
|
|
self.subprocess: asyncio.subprocess.Process | None = None
|
|
self.stdout_buffer: list[str] = []
|
|
self.stderr_buffer: list[str] = []
|
|
|
|
async def start(self):
|
|
project_root = Path(__file__).parent.parent.parent
|
|
runner_path = project_root / "src" / "main.py"
|
|
|
|
env_vars = os.environ.copy()
|
|
env_vars[ENV_GRANT_TOKEN] = "test_token"
|
|
env_vars[ENV_TASK_BROKER_URI] = self.task_broker_url
|
|
env_vars[ENV_TASK_TIMEOUT] = str(TASK_TIMEOUT)
|
|
env_vars[ENV_HEALTH_CHECK_SERVER_ENABLED] = "true"
|
|
env_vars[ENV_HEALTH_CHECK_SERVER_PORT] = str(HEALTH_CHECK_PORT)
|
|
if self.graceful_shutdown_timeout is not None:
|
|
env_vars[ENV_GRACEFUL_SHUTDOWN_TIMEOUT] = str(
|
|
self.graceful_shutdown_timeout
|
|
)
|
|
env_vars["PYTHONPATH"] = str(project_root)
|
|
|
|
self.subprocess = await asyncio.create_subprocess_exec(
|
|
sys.executable,
|
|
str(runner_path),
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
env=env_vars,
|
|
cwd=str(project_root),
|
|
)
|
|
|
|
asyncio.create_task(self._read_stdout())
|
|
asyncio.create_task(self._read_stderr())
|
|
|
|
def is_running(self) -> bool:
|
|
return self.subprocess is not None and self.subprocess.returncode is None
|
|
|
|
async def stop(self) -> None:
|
|
if not self.subprocess or self.subprocess.returncode is not None:
|
|
return
|
|
|
|
self.subprocess.terminate()
|
|
try:
|
|
await asyncio.wait_for(
|
|
self.subprocess.wait(), timeout=GRACEFUL_SHUTDOWN_TIMEOUT
|
|
)
|
|
except asyncio.TimeoutError:
|
|
self.subprocess.kill()
|
|
await self.subprocess.wait()
|
|
|
|
async def _read_stdout(self):
|
|
if not self.subprocess or not self.subprocess.stdout:
|
|
return
|
|
|
|
while True:
|
|
line = await self.subprocess.stdout.readline()
|
|
if not line:
|
|
break
|
|
self.stdout_buffer.append(line.decode().strip())
|
|
|
|
async def _read_stderr(self):
|
|
if not self.subprocess or not self.subprocess.stderr:
|
|
return
|
|
|
|
while True:
|
|
line = await self.subprocess.stderr.readline()
|
|
if not line:
|
|
break
|
|
self.stderr_buffer.append(line.decode().strip())
|