Fix mtsender enqueuer and complete test

This commit is contained in:
Lonami Exo 2023-08-31 12:27:43 +02:00
parent dfc540c472
commit 7166059132
2 changed files with 30 additions and 14 deletions

View File

@ -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,

View File

@ -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())