from __future__ import annotations

import asyncio
import unittest
from unittest.mock import patch

from routes import voice_webrtc


class FakeConnection:
    pc_id = "pc_test"


class WebRTCPipelineLifecycleTest(unittest.IsolatedAsyncioTestCase):
    async def asyncTearDown(self):
        tasks = list(voice_webrtc._pipeline_tasks)
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        voice_webrtc._pipeline_tasks.clear()
        voice_webrtc._connection_tasks.clear()
        voice_webrtc._active_connections.clear()

    async def test_pipeline_task_is_retained_until_completion(self):
        connection = FakeConnection()
        release = asyncio.Event()

        async def pipeline():
            await release.wait()

        task = voice_webrtc._start_pipeline_task(connection, pipeline())
        self.assertIn(task, voice_webrtc._pipeline_tasks)
        self.assertIs(voice_webrtc._connection_tasks[connection], task)

        release.set()
        await task
        await asyncio.sleep(0)

        self.assertNotIn(task, voice_webrtc._pipeline_tasks)
        self.assertNotIn(connection, voice_webrtc._connection_tasks)

    async def test_stuck_pipeline_is_cancelled_after_close_grace_period(self):
        connection = FakeConnection()
        cancelled = asyncio.Event()

        async def pipeline():
            try:
                await asyncio.Event().wait()
            finally:
                cancelled.set()

        task = voice_webrtc._start_pipeline_task(connection, pipeline())
        with patch.object(voice_webrtc, "PIPELINE_CLOSE_GRACE_SECONDS", 0.01):
            await voice_webrtc._wait_for_pipeline_close(
                task,
                connection_id=connection.pc_id,
            )

        self.assertTrue(cancelled.is_set())
        self.assertTrue(task.cancelled())


if __name__ == "__main__":
    unittest.main()