LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2025-02-13 18:11:06 +03:00
Code Issues Packages Projects Releases Wiki Activity

Fix mtsender enqueuer and complete test

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
client/src/telethon/_impl/mtsender/sender.py
View File

@ -71,11 +71,8 @@ class Request(Generic[Return]):
class Enqueuer: class Enqueuer:
__slots__ = ("_queue",) __slots__ = ("_queue",)
def __init__( def __init__(self, queue: Queue[Request[object]]) -> None:
self, self._queue = queue
) -> None:
# TODO use a bound
self._queue: Queue[Request[object]] = Queue()
def enqueue(self, request: RemoteCall[Return]) -> Future[Return]: def enqueue(self, request: RemoteCall[Return]) -> Future[Return]:
body = bytes(request) body = bytes(request)
@ -95,7 +92,6 @@ class Sender:
_mtp: Mtp _mtp: Mtp
_mtp_buffer: bytearray _mtp_buffer: bytearray
_requests: List[Request[object]] _requests: List[Request[object]]
_request_tx: Queue[Request[object]]
_request_rx: Queue[Request[object]] _request_rx: Queue[Request[object]]
_next_ping: float _next_ping: float
_read_buffer: bytearray _read_buffer: bytearray
@ -106,8 +102,7 @@ class Sender:
cls, transport: Transport, mtp: Mtp, addr: str cls, transport: Transport, mtp: Mtp, addr: str
) -> Tuple[Self, Enqueuer]: ) -> Tuple[Self, Enqueuer]:
reader, writer = await asyncio.open_connection(*addr.split(":")) reader, writer = await asyncio.open_connection(*addr.split(":"))
tx: Queue[object] = Queue() request_queue: Queue[object] = Queue()
rx = tx
return ( return (
cls( cls(
@ -117,13 +112,12 @@ class Sender:
_mtp=mtp, _mtp=mtp,
_mtp_buffer=bytearray(), _mtp_buffer=bytearray(),
_requests=[], _requests=[],
_request_tx=tx, _request_rx=request_queue,
_request_rx=rx,
_next_ping=asyncio.get_running_loop().time() + PING_DELAY, _next_ping=asyncio.get_running_loop().time() + PING_DELAY,
_read_buffer=bytearray(), _read_buffer=bytearray(),
_write_drain_pending=False, _write_drain_pending=False,
), ),
Enqueuer(), Enqueuer(request_queue),
) )
async def invoke(self, request: RemoteCall[Return]) -> bytes: 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=Encrypted(auth_key, time_offset=time_offset, first_salt=first_salt),
_mtp_buffer=sender._mtp_buffer, _mtp_buffer=sender._mtp_buffer,
_requests=sender._requests, _requests=sender._requests,
_request_tx=sender._request_tx,
_request_rx=sender._request_rx, _request_rx=sender._request_rx,
_next_ping=time.time() + PING_DELAY, _next_ping=time.time() + PING_DELAY,
_read_buffer=sender._read_buffer, _read_buffer=sender._read_buffer,

27
client/tests/mtsender_test.py
View File

@ -4,6 +4,7 @@ import logging
from pytest import LogCaptureFixture from pytest import LogCaptureFixture
from telethon._impl.mtproto.transport.full import Full from telethon._impl.mtproto.transport.full import Full
from telethon._impl.mtsender.sender import connect 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" 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() connect(Full(), TELEGRAM_DEFAULT_TEST_DC), timeout()
) )
# TODO test enqueuer rx = enqueuer.enqueue(
sender, enqueuer 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()) asyncio.run(func())