CJ_Clippy
de5a4c11b2
subrepo: subdir: "apps/twscrape" merged: "e4dce5d3" upstream: origin: "https://github.com/vladkens/twscrape" branch: "main" commit: "e4dce5d3" git-subrepo: version: "0.4.9" origin: "???" commit: "???"
159 lines
4.8 KiB
Python
159 lines
4.8 KiB
Python
from contextlib import aclosing
|
|
|
|
import httpx
|
|
from pytest_httpx import HTTPXMock
|
|
|
|
from twscrape.accounts_pool import AccountsPool
|
|
from twscrape.queue_client import QueueClient
|
|
|
|
DB_FILE = "/tmp/twscrape_test_queue_client.db"
|
|
URL = "https://example.com/api"
|
|
CF = tuple[AccountsPool, QueueClient]
|
|
|
|
|
|
async def get_locked(pool: AccountsPool) -> set[str]:
|
|
rep = await pool.get_all()
|
|
return set([x.username for x in rep if x.locks.get("SearchTimeline", None) is not None])
|
|
|
|
|
|
async def test_lock_account_when_used(httpx_mock: HTTPXMock, client_fixture):
|
|
pool, client = client_fixture
|
|
|
|
locked = await get_locked(pool)
|
|
assert len(locked) == 0
|
|
|
|
# should lock account on getting it
|
|
await client.__aenter__()
|
|
locked = await get_locked(pool)
|
|
assert len(locked) == 1
|
|
assert "user1" in locked
|
|
|
|
# keep locked on request
|
|
httpx_mock.add_response(url=URL, json={"foo": "bar"}, status_code=200)
|
|
assert (await client.get(URL)).json() == {"foo": "bar"}
|
|
|
|
locked = await get_locked(pool)
|
|
assert len(locked) == 1
|
|
assert "user1" in locked
|
|
|
|
# unlock on exit
|
|
await client.__aexit__(None, None, None)
|
|
locked = await get_locked(pool)
|
|
assert len(locked) == 0
|
|
|
|
|
|
async def test_do_not_switch_account_on_200(httpx_mock: HTTPXMock, client_fixture: CF):
|
|
pool, client = client_fixture
|
|
|
|
# get account and lock it
|
|
await client.__aenter__()
|
|
locked1 = await get_locked(pool)
|
|
assert len(locked1) == 1
|
|
|
|
# make several requests with status=200
|
|
for x in range(1):
|
|
httpx_mock.add_response(url=URL, json={"foo": x}, status_code=200)
|
|
rep = await client.get(URL)
|
|
assert rep is not None
|
|
assert rep.json() == {"foo": x}
|
|
|
|
# account should not be switched
|
|
locked2 = await get_locked(pool)
|
|
assert locked1 == locked2
|
|
|
|
# unlock on exit
|
|
await client.__aexit__(None, None, None)
|
|
locked3 = await get_locked(pool)
|
|
assert len(locked3) == 0
|
|
|
|
|
|
async def test_switch_acc_on_http_error(httpx_mock: HTTPXMock, client_fixture: CF):
|
|
pool, client = client_fixture
|
|
|
|
# locked account on enter
|
|
await client.__aenter__()
|
|
locked1 = await get_locked(pool)
|
|
assert len(locked1) == 1
|
|
|
|
# fail one request, account should be switched
|
|
httpx_mock.add_response(url=URL, json={"foo": "1"}, status_code=403)
|
|
httpx_mock.add_response(url=URL, json={"foo": "2"}, status_code=200)
|
|
|
|
rep = await client.get(URL)
|
|
assert rep is not None
|
|
assert rep.json() == {"foo": "2"}
|
|
|
|
locked2 = await get_locked(pool)
|
|
assert len(locked2) == 2
|
|
|
|
# unlock on exit (failed account still should locked)
|
|
await client.__aexit__(None, None, None)
|
|
locked3 = await get_locked(pool)
|
|
assert len(locked3) == 1
|
|
assert locked1 == locked3 # failed account locked
|
|
|
|
|
|
async def test_retry_with_same_acc_on_network_error(httpx_mock: HTTPXMock, client_fixture: CF):
|
|
pool, client = client_fixture
|
|
|
|
# locked account on enter
|
|
await client.__aenter__()
|
|
locked1 = await get_locked(pool)
|
|
assert len(locked1) == 1
|
|
|
|
# timeout on first request, account should not be switched
|
|
httpx_mock.add_exception(httpx.ReadTimeout("Unable to read within timeout"))
|
|
httpx_mock.add_response(url=URL, json={"foo": "2"}, status_code=200)
|
|
|
|
rep = await client.get(URL)
|
|
assert rep is not None
|
|
assert rep.json() == {"foo": "2"}
|
|
|
|
locked2 = await get_locked(pool)
|
|
assert locked2 == locked1
|
|
|
|
# check username added to request obj (for logging)
|
|
username = getattr(rep, "__username", None)
|
|
assert username is not None
|
|
|
|
|
|
async def test_ctx_closed_on_break(httpx_mock: HTTPXMock, client_fixture: CF):
|
|
pool, client = client_fixture
|
|
|
|
async def get_data_stream():
|
|
async with client as c:
|
|
counter = 0
|
|
while True:
|
|
counter += 1
|
|
check_retry = counter == 2
|
|
before_ctx = c.ctx
|
|
|
|
if check_retry:
|
|
httpx_mock.add_response(url=URL, json={"counter": counter}, status_code=403)
|
|
httpx_mock.add_response(url=URL, json={"counter": counter}, status_code=200)
|
|
else:
|
|
httpx_mock.add_response(url=URL, json={"counter": counter}, status_code=200)
|
|
|
|
rep = await c.get(URL)
|
|
|
|
if check_retry:
|
|
assert before_ctx != c.ctx
|
|
elif before_ctx is not None:
|
|
assert before_ctx == c.ctx
|
|
|
|
assert rep is not None
|
|
assert rep.json() == {"counter": counter}
|
|
yield rep.json()["counter"]
|
|
|
|
if counter == 9:
|
|
return
|
|
|
|
# need to use async with to break to work
|
|
async with aclosing(get_data_stream()) as gen:
|
|
async for x in gen:
|
|
if x == 3:
|
|
break
|
|
|
|
# ctx should be None after break
|
|
assert client.ctx is None
|