ai-video-fullstack/backend/tests/test_background_task_native.py

"""Real Pipecat scheduling with an explicitly simulated remote executor."""
import asyncio
import json
import unittest
from pathlib import Path

from pipecat.frames.frames import CancelFrame
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.worker import PipelineWorker
from pipecat.processors.frame_processor import FrameProcessor
from services.background_tasks.native import NativeTaskRuntime


class Passthrough(FrameProcessor):
    async def process_frame(self, frame, direction):
        await super().process_frame(frame, direction)
        await self.push_frame(frame, direction)


class NativeJobTest(unittest.IsolatedAsyncioTestCase):
    async def test_native_timeout_calls_remote_cleanup(self):
        cancelled = asyncio.Event()

        async def execute(message, worker):
            await asyncio.Event().wait()

        async def cleanup(message):
            cancelled.set()

        async def ignore(message):
            pass

        runtime = NativeTaskRuntime(execute, cleanup, ignore, ignore)
        await runtime.start()
        try:
            await runtime.submit("timeout-test", 0.05)
            await asyncio.wait_for(cancelled.wait(), 2)
        finally:
            await runtime.close()

    async def test_progress_completion_and_media_disconnect(self):
        progress, finished, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
        results = []

        async def execute(message, worker):
            # SIMULATED executor; no remote OpenCode or model is called here.
            await worker.send_job_update(message.job_id, {"step": "working"})
            await release.wait()
            return {"result": "done"}

        async def cancelled(message):
            pass

        async def update(message):
            progress.set()

        async def result(message):
            results.append(message)
            finished.set()

        runtime = NativeTaskRuntime(execute, cancelled, update, result)
        await runtime.start()
        try:
            job_id = await runtime.submit("test-task", 5)
            await asyncio.wait_for(progress.wait(), 2)
            self.assertFalse(finished.is_set())
            media = PipelineWorker(Pipeline([Passthrough()]), enable_rtvi=False)
            await runtime.runner.add_workers(media)
            await media.queue_frame(CancelFrame(reason="client_disconnected"))
            await asyncio.wait_for(media.wait(), 2)
            self.assertIn(job_id, runtime.executor.active_jobs)
            release.set()
            await asyncio.wait_for(finished.wait(), 2)
            self.assertEqual(results[0].status, "completed")
            self.assertEqual(results[0].response, {"result": "done"})
            artifact = Path(__file__).resolve().parents[2] / "implement-stage/native-jobs-smoke.json"
            artifact.write_text(json.dumps({"simulated_executor": True, "native_job": "completed", "survives_media_disconnect": True}) + "\n")
        finally:
            await runtime.close()