pengjunjie/langflow
1
0
Fork 1
mirror of https://github.com/langflow-ai/langflow.git synced 2026-07-24 10:50:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

test(background): restore manager DB service + settings on subprocess harness teardown to prevent state leak

Browse Source
This commit is contained in:
ogabrielluiz
2026-06-04 07:26:44 -03:00
parent e2be9f932a
commit 20f5c6940f
1 changed files with 23 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
src/backend/tests/unit/services/background_execution/_subprocess_harness.py
View File

@ -64,6 +64,9 @@ class WorkerHarness:
_client: object | None = None
_prior_db_env: str | None = None
_prior_db_env_set: bool = False
_prior_db_service: object | None = None
_prior_settings_url: str | None = None
_db_service: object | None = None
async def seed_real_flow(self, *, input_value: str = "hello") -> tuple[uuid.UUID, uuid.UUID]:
"""Seed a real user + a real ChatInput->ChatOutput flow; return (user_id, flow_id).
@ -213,7 +216,21 @@ class WorkerHarness:
await self._client.flushdb()
with contextlib.suppress(Exception):
await self._client.aclose()
# Restore the LANGFLOW_DATABASE_URL env so it does not leak to later tests.
# Tear down the pg engine and restore the manager's DB service + settings
# url + env var so this harness does not leak state to later tests.
from langflow.services.deps import get_settings_service
from lfx.services.manager import get_service_manager
from lfx.services.schema import ServiceType
manager = get_service_manager()
manager.services.pop(ServiceType.DATABASE_SERVICE, None)
if self._db_service is not None:
with contextlib.suppress(Exception):
await self._db_service.teardown()
if self._prior_db_service is not None:
manager.services[ServiceType.DATABASE_SERVICE] = self._prior_db_service
if self._prior_settings_url is not None:
get_settings_service().settings.database_url = self._prior_settings_url
if self._prior_db_env_set:
os.environ["LANGFLOW_DATABASE_URL"] = self._prior_db_env # type: ignore[assignment]
else:
@ -245,9 +262,10 @@ async def setup_worker_harness(pg_uri: str, redis_base_url: str, *, redis_db: in
os.environ["LANGFLOW_DATABASE_URL"] = db_url
settings_service = get_settings_service()
prior_settings_url = settings_service.settings.database_url
settings_service.settings.database_url = db_url
manager = get_service_manager()
manager.services.pop(ServiceType.DATABASE_SERVICE, None)
prior_db_service = manager.services.pop(ServiceType.DATABASE_SERVICE, None)
db_service = DatabaseServiceFactory().create(settings_service)
manager.services[ServiceType.DATABASE_SERVICE] = db_service
await db_service.create_db_and_tables()
@ -258,6 +276,9 @@ async def setup_worker_harness(pg_uri: str, redis_base_url: str, *, redis_db: in
job_service=JobService(),
_prior_db_env=prior_db_env,
_prior_db_env_set=prior_db_env_set,
_prior_db_service=prior_db_service,
_prior_settings_url=prior_settings_url,
_db_service=db_service,
)
# Flush the dedicated redis DB so a prior run cannot leak claim-queue ids.
await harness.client.flushdb()