Repository: gradio-app/gradio
Pull request: Tests for Queue
Adds unit tests for queuing. In some cases, it was hard to test behavior and not implementation, so any feedback on making the tests better is welcome. But I think it's important to have some unit tests for queuing in place as we work on improving it to make sure we don't inadvertently break anything.

Fixes: #1983

Some additional notes:
* `event_queue.py` has been renamed to `queue.py`

Problem statement:
Write Tests for Queueing
**Is your feature request related to a problem? Please describe.**  
We need to write tests for the new event queue in the Frontend.

**Additional context**  
Websocket testing tooling is seems to not work well in Python, and we're unable to test it in Backend.

Testing a single event case and multiple events case would be nice.

Hints from issue discussion:
We should write unit tests for `queueing.py` as well!
@abidlabs I'd like to write unit tests for 'queueing.py' if it is still unassigned. Please let me know
Hi @hibestil, that would be a welcome contribution, thank you!

Commit message: removed websocket

Patch to review:
diff --git a/gradio/blocks.py b/gradio/blocks.py
index b6d0fe9bd6..67b722dd47 100644
--- a/gradio/blocks.py
+++ b/gradio/blocks.py
@@ -20,9 +20,9 @@
 from gradio import (
     components,
     encryptor,
-    event_queue,
     external,
     networking,
+    queue,
     routes,
     strings,
     utils,
@@ -910,7 +910,7 @@ def queue(
             demo.launch()
         """
         self.enable_queue = default_enabled
-        self._queue = event_queue.Queue(
+        self._queue = queue.Queue(
             live_updates=status_update_rate == "auto",
             concurrency_count=concurrency_count,
             data_gathering_start=client_position_to_load_data,
diff --git a/gradio/event_queue.py b/gradio/queue.py
similarity index 99%
rename from gradio/event_queue.py
rename to gradio/queue.py
index 639b0ca8de..018910a49d 100644
--- a/gradio/event_queue.py
+++ b/gradio/queue.py
@@ -184,6 +184,7 @@ async def send_estimation(
         client_awake = await self.send_message(event, estimation.dict())
         if not client_awake:
             await self.clean_event(event)
+        return estimation
 
     def update_estimation(self, duration: float) -> None:
         """
diff --git a/gradio/routes.py b/gradio/routes.py
index 4637f497d9..b10c9d581c 100644
--- a/gradio/routes.py
+++ b/gradio/routes.py
@@ -29,8 +29,8 @@
 
 import gradio
 from gradio import encryptor
-from gradio.event_queue import Estimation, Event, Queue
 from gradio.exceptions import Error
+from gradio.queue import Estimation, Event, Queue
 
 mimetypes.init()
 
diff --git a/test/test_event_queue.py b/test/test_event_queue.py
deleted file mode 100644
index 2dce5505bc..0000000000
--- a/test/test_event_queue.py
+++ /dev/null
@@ -1,13 +0,0 @@
-import asyncio
-import os
-
-import pytest
-from fastapi.testclient import TestClient
-
-import gradio as gr
-
-os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
-
-
-class TestQueue:
-    pass  # TODO
diff --git a/test/test_queue.py b/test/test_queue.py
new file mode 100644
index 0000000000..654c36c1bf
--- /dev/null
+++ b/test/test_queue.py
@@ -0,0 +1,155 @@
+import os
+from unittest.mock import MagicMock
+
+import pytest
+
+from gradio.queue import Event, Queue
+
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+
+
+class AsyncMock(MagicMock):
+    async def __call__(self, *args, **kwargs):
+        return super(AsyncMock, self).__call__(*args, **kwargs)
+
+
+@pytest.fixture()
+def queue() -> Queue:
+    queue_object = Queue(
+        live_updates=True,
+        concurrency_count=1,
+        data_gathering_start=1,
+        update_intervals=1,
+        max_size=None,
+    )
+    yield queue_object
+    queue_object.close()
+
+
+@pytest.fixture()
+def mock_event() -> Event:
+    websocket = MagicMock()
+    event = Event(websocket=websocket)
+    yield event
+
+
+class TestQueueMethods:
+    @pytest.mark.asyncio
+    async def test_start(self, queue: Queue):
+        await queue.start()
+        assert queue.stopped is False
+        assert queue.get_active_worker_count() == 0
+
+    @pytest.mark.asyncio
+    async def test_stop_resume(self, queue: Queue):
+        await queue.start()
+        queue.close()
+        assert queue.stopped
+        queue.resume()
+        assert queue.stopped is False
+
+    @pytest.mark.asyncio
+    async def test_receive(self, queue: Queue, mock_event: Event):
+        await queue.get_message(mock_event)
+        assert mock_event.websocket.receive_json.called
+
+    @pytest.mark.asyncio
+    async def test_send(self, queue: Queue, mock_event: Event):
+        await queue.send_message(mock_event, {})
+        assert mock_event.websocket.send_json.called
+
+    @pytest.mark.asyncio
+    async def test_add_to_queue(self, queue: Queue, mock_event: Event):
+        queue.push(mock_event)
+        assert len(queue.event_queue) == 1
+
+    @pytest.mark.asyncio
+    async def test_add_to_queue_with_max_size(self, queue: Queue, mock_event: Event):
+        queue.max_size = 1
+        queue.push(mock_event)
+        assert len(queue.event_queue) == 1
+        queue.push(mock_event)
+        assert len(queue.event_queue) == 1
+
+    @pytest.mark.asyncio
+    async def test_clean_event(self, queue: Queue, mock_event: Event):
+        queue.push(mock_event)
+        await queue.clean_event(mock_event)
+        assert len(queue.event_queue) == 0
+
+    @pytest.mark.asyncio
+    async def test_gather_event_data(self, queue: Queue, mock_event: Event):
+        queue.send_message = AsyncMock()
+        queue.get_message = AsyncMock()
+        queue.send_message.return_value = True
+        queue.get_message.return_value = {"data": ["test"], "fn": 0}
+
+        assert await queue.gather_event_data(mock_event)
+        assert queue.send_message.called
+        assert mock_event.data == {"data": ["test"], "fn": 0}
+
+        queue.send_message.called = False
+        assert await queue.gather_event_data(mock_event)
+        assert not (queue.send_message.called)
+
+    @pytest.mark.asyncio
+    async def test_gather_data_for_first_ranks(self, queue: Queue, mock_event: Event):
+        websocket = MagicMock()
+        mock_event2 = Event(websocket=websocket)
+        queue.send_message = AsyncMock()
+        queue.get_message = AsyncMock()
+        queue.send_message.return_value = True
+        queue.get_message.return_value = {"data": ["test"], "fn": 0}
+
+        queue.push(mock_event)
+        queue.push(mock_event2)
+        await queue.gather_data_for_first_ranks()
+        assert mock_event.data is not None
+        assert mock_event2.data is None
+
+
+class TestQueueEstimation:
+    def test_get_update_estimation(self, queue: Queue):
+        queue.update_estimation(5)
+        estimation = queue.get_estimation()
+        assert estimation.avg_event_process_time == 5
+
+        queue.update_estimation(15)
+        estimation = queue.get_estimation()
+        assert estimation.avg_event_process_time == 10  # (5 + 15) / 2
+
+        queue.update_estimation(100)
+        estimation = queue.get_estimation()
+        assert estimation.avg_event_process_time == 40  # (5 + 15 + 100) / 3
+
+    @pytest.mark.asyncio
+    async def test_send_estimation(self, queue: Queue, mock_event: Event):
+        queue.send_message = AsyncMock()
+        queue.send_message.return_value = True
+        estimation = queue.get_estimation()
+        estimation = await queue.send_estimation(mock_event, estimation, 1)
+        assert queue.send_message.called
+        assert estimation.rank == 1
+
+        queue.update_estimation(5)
+        estimation = await queue.send_estimation(mock_event, estimation, 2)
+        assert estimation.rank == 2
+        assert estimation.rank_eta == 15
+
+
+class TestQueueProcessEvents:
+    @pytest.mark.asyncio
+    async def test_process_event(self, queue: Queue, mock_event: Event):
+        queue.gather_event_data = AsyncMock()
+        queue.gather_event_data.return_value = True
+        queue.send_message = AsyncMock()
+        queue.send_message.return_value = True
+        queue.call_prediction = AsyncMock()
+        queue.call_prediction.return_value = MagicMock()
+        queue.call_prediction.return_value.json = {"is_generating": False}
+        mock_event.disconnect = AsyncMock()
+        queue.clean_event = AsyncMock()
+        await queue.process_event(mock_event)
+
+        assert queue.call_prediction.called
+        assert mock_event.disconnect.called
diff --git a/test/test_routes.py b/test/test_routes.py
index 377157fe47..f27b2e4a88 100644
--- a/test/test_routes.py
+++ b/test/test_routes.py
@@ -1,9 +1,11 @@
 """Contains tests for networking.py and app.py"""
 
+import asyncio
 import os
 import unittest
-import unittest.mock as mock
+from unittest.mock import MagicMock
 
+import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
@@ -12,6 +14,11 @@
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 
 
+class AsyncMock(MagicMock):
+    async def __call__(self, *args, **kwargs):
+        return super(AsyncMock, self).__call__(*args, **kwargs)
+
+
 class TestRoutes(unittest.TestCase):
     def setUp(self) -> None:
         self.io = Interface(lambda x: x + x, "text", "text")
@@ -22,10 +29,6 @@ def test_get_main_route(self):
         response = self.client.get("/")
         self.assertEqual(response.status_code, 200)
 
-    # def test_get_api_route(self):
-    #     response = self.client.get("/api/")
-    #     self.assertEqual(response.status_code, 200)
-
     def test_static_files_served_safely(self):
         # Make sure things outside the static folder are not accessible
         response = self.client.get(r"/static/..%2findex.html")


Code review comments:
diff_hunk: @@ -12,6 +14,11 @@
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 
 
+class AsyncMock(MagicMock):
path: test/test_routes.py
text: @user1:
This class isn't used anywhere