From 71660591329549b082ee7dd09845cbda12d62f97 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 31 Aug 2023 12:27:43 +0200
Subject: [PATCH] Fix mtsender enqueuer and complete test

---
 client/src/telethon/_impl/mtsender/sender.py | 17 ++++--------
 client/tests/mtsender_test.py                | 27 ++++++++++++++++++--
 2 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/client/src/telethon/_impl/mtsender/sender.py b/client/src/telethon/_impl/mtsender/sender.py
index eb9facd3..98817b7a 100644
--- a/client/src/telethon/_impl/mtsender/sender.py
+++ b/client/src/telethon/_impl/mtsender/sender.py
@@ -71,11 +71,8 @@ class Request(Generic[Return]):
 class Enqueuer:
     __slots__ = ("_queue",)
 
-    def __init__(
-        self,
-    ) -> None:
-        # TODO use a bound
-        self._queue: Queue[Request[object]] = Queue()
+    def __init__(self, queue: Queue[Request[object]]) -> None:
+        self._queue = queue
 
     def enqueue(self, request: RemoteCall[Return]) -> Future[Return]:
         body = bytes(request)
@@ -95,7 +92,6 @@ class Sender:
     _mtp: Mtp
     _mtp_buffer: bytearray
     _requests: List[Request[object]]
-    _request_tx: Queue[Request[object]]
     _request_rx: Queue[Request[object]]
     _next_ping: float
     _read_buffer: bytearray
@@ -106,8 +102,7 @@ class Sender:
         cls, transport: Transport, mtp: Mtp, addr: str
     ) -> Tuple[Self, Enqueuer]:
         reader, writer = await asyncio.open_connection(*addr.split(":"))
-        tx: Queue[object] = Queue()
-        rx = tx
+        request_queue: Queue[object] = Queue()
 
         return (
             cls(
@@ -117,13 +112,12 @@ class Sender:
                 _mtp=mtp,
                 _mtp_buffer=bytearray(),
                 _requests=[],
-                _request_tx=tx,
-                _request_rx=rx,
+                _request_rx=request_queue,
                 _next_ping=asyncio.get_running_loop().time() + PING_DELAY,
                 _read_buffer=bytearray(),
                 _write_drain_pending=False,
             ),
-            Enqueuer(),
+            Enqueuer(request_queue),
         )
 
     async def invoke(self, request: RemoteCall[Return]) -> bytes:
@@ -315,7 +309,6 @@ async def generate_auth_key(
             _mtp=Encrypted(auth_key, time_offset=time_offset, first_salt=first_salt),
             _mtp_buffer=sender._mtp_buffer,
             _requests=sender._requests,
-            _request_tx=sender._request_tx,
             _request_rx=sender._request_rx,
             _next_ping=time.time() + PING_DELAY,
             _read_buffer=sender._read_buffer,
diff --git a/client/tests/mtsender_test.py b/client/tests/mtsender_test.py
index b187fb7f..a23cfcb8 100644
--- a/client/tests/mtsender_test.py
+++ b/client/tests/mtsender_test.py
@@ -4,6 +4,7 @@ import logging
 from pytest import LogCaptureFixture
 from telethon._impl.mtproto.transport.full import Full
 from telethon._impl.mtsender.sender import connect
+from telethon._impl.tl import LAYER, abcs, functions, types
 
 TELEGRAM_TEST_DC_2 = "149.154.167.40:443"
 
@@ -25,7 +26,29 @@ def test_invoke_encrypted_method(caplog: LogCaptureFixture) -> None:
             connect(Full(), TELEGRAM_DEFAULT_TEST_DC), timeout()
         )
 
-        # TODO test enqueuer
-        sender, enqueuer
+        rx = enqueuer.enqueue(
+            functions.invoke_with_layer(
+                layer=LAYER,
+                query=functions.init_connection(
+                    api_id=1,
+                    device_model="Test",
+                    system_version="0.1",
+                    app_version="0.1",
+                    system_lang_code="en",
+                    lang_pack="",
+                    lang_code="",
+                    proxy=None,
+                    params=None,
+                    query=functions.help.get_nearest_dc(),
+                ),
+            )
+        )
+
+        while True:
+            await asyncio.wait_for(sender.step(), timeout=timeout())
+            if rx.done():
+                nearest = abcs.NearestDc.from_bytes(rx.result())
+                assert isinstance(nearest, types.NearestDc)
+                break
 
     asyncio.run(func())