mirror of
https://github.com/vladkens/twscrape.git
synced 2026-10-16 03:04:00 -04:00
383 lines
13 KiB
Python
383 lines
13 KiB
Python
import pytest
|
|
|
|
from twscrape.accounts_pool import AccountsPool, NoAccountError
|
|
from twscrape.api import API
|
|
from twscrape.utils import utc
|
|
|
|
|
|
async def test_add_accounts(pool_mock: AccountsPool):
|
|
# should add account
|
|
await pool_mock.add_account("user1", "pass1", "email1", "email_pass1")
|
|
acc = await pool_mock.get("user1")
|
|
assert acc.username == "user1"
|
|
assert acc.password == "pass1"
|
|
assert acc.email == "email1"
|
|
assert acc.email_password == "email_pass1"
|
|
acc.email_password = "_"
|
|
assert acc.login_method == "password"
|
|
|
|
# should not add account with same username
|
|
await pool_mock.add_account("user1", "pass2", "email2", "email_pass2")
|
|
acc = await pool_mock.get("user1")
|
|
assert acc.username == "user1"
|
|
assert acc.password == "pass1"
|
|
assert acc.email == "email1"
|
|
assert acc.email_password == "email_pass1"
|
|
|
|
# should not add account with different username case
|
|
await pool_mock.add_account("USER1", "pass2", "email2", "email_pass2")
|
|
acc = await pool_mock.get("user1")
|
|
assert acc.username == "user1"
|
|
assert acc.password == "pass1"
|
|
assert acc.email == "email1"
|
|
assert acc.email_password == "email_pass1"
|
|
|
|
# should add account with different username
|
|
await pool_mock.add_account("user2", "pass2", "email2", "email_pass2")
|
|
acc = await pool_mock.get("user2")
|
|
assert acc.username == "user2"
|
|
assert acc.password == "pass2"
|
|
assert acc.email == "email2"
|
|
assert acc.email_password == "email_pass2"
|
|
|
|
|
|
async def test_add_account_cookies(pool_mock: AccountsPool):
|
|
await pool_mock.add_account_cookies("user1", "auth_token=token; ct0=csrf")
|
|
acc = await pool_mock.get("user1")
|
|
|
|
assert acc.active is True
|
|
assert acc.has_session is True
|
|
assert acc.login_method == "cookies"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"cookies",
|
|
["auth_token=token", "ct0=csrf", "other=value"],
|
|
)
|
|
async def test_add_account_cookies_requires_full_session(pool_mock: AccountsPool, cookies: str):
|
|
with pytest.raises(ValueError):
|
|
await pool_mock.add_account_cookies("user1", cookies)
|
|
assert await pool_mock.get_account("user1") is None
|
|
|
|
|
|
async def test_add_account_cookies_existing(pool_mock: AccountsPool):
|
|
await pool_mock.add_account(
|
|
"user1",
|
|
"pass",
|
|
"email",
|
|
"email-pass",
|
|
proxy="http://proxy.test",
|
|
mfa_code="mfa-secret",
|
|
cookies="auth_token=token; ct0=csrf",
|
|
)
|
|
acc = await pool_mock.get("user1")
|
|
acc.stats = {"SearchTimeline": 3}
|
|
await pool_mock.save(acc)
|
|
|
|
await pool_mock.add_account_cookies("user1", "auth_token=other; ct0=other-csrf")
|
|
same = await pool_mock.get("user1")
|
|
|
|
assert same.cookies == {"auth_token": "other", "ct0": "other-csrf"}
|
|
assert same.active is True
|
|
assert same.login_method == "password"
|
|
assert same.password == "pass"
|
|
assert same.proxy == acc.proxy
|
|
assert same.stats == acc.stats
|
|
|
|
|
|
async def test_get_all(pool_mock: AccountsPool):
|
|
# should return empty list
|
|
accs = await pool_mock.get_all()
|
|
assert len(accs) == 0
|
|
|
|
# should return all accounts
|
|
await pool_mock.add_account("user1", "pass1", "email1", "email_pass1")
|
|
await pool_mock.add_account("user2", "pass2", "email2", "email_pass2")
|
|
accs = await pool_mock.get_all()
|
|
assert len(accs) == 2
|
|
assert accs[0].username == "user1"
|
|
assert accs[1].username == "user2"
|
|
|
|
|
|
async def test_save(pool_mock: AccountsPool):
|
|
# should save account
|
|
await pool_mock.add_account("user1", "pass1", "email1", "email_pass1")
|
|
acc = await pool_mock.get("user1")
|
|
acc.password = "pass2"
|
|
await pool_mock.save(acc)
|
|
acc = await pool_mock.get("user1")
|
|
assert acc.password == "pass2"
|
|
|
|
# should not save account
|
|
acc = await pool_mock.get("user1")
|
|
acc.username = "user2"
|
|
await pool_mock.save(acc)
|
|
acc = await pool_mock.get("user1")
|
|
assert acc.username == "user1"
|
|
|
|
|
|
async def test_get_for_queue(pool_mock: AccountsPool):
|
|
Q = "test_queue"
|
|
|
|
# should return account
|
|
await pool_mock.add_account("user1", "pass1", "email1", "email_pass1")
|
|
await pool_mock.set_active("user1", True)
|
|
acc = await pool_mock.get_for_queue(Q)
|
|
assert acc is not None
|
|
assert acc.username == "user1"
|
|
assert acc.active is True
|
|
assert acc.locks is not None
|
|
assert Q in acc.locks
|
|
assert acc.locks[Q] is not None
|
|
|
|
# should return None
|
|
acc = await pool_mock.get_for_queue(Q)
|
|
assert acc is None
|
|
|
|
|
|
async def test_account_unlock(pool_mock: AccountsPool):
|
|
Q = "test_queue"
|
|
|
|
await pool_mock.add_account("user1", "pass1", "email1", "email_pass1")
|
|
await pool_mock.set_active("user1", True)
|
|
acc = await pool_mock.get_for_queue(Q)
|
|
assert acc is not None
|
|
assert acc.locks[Q] is not None
|
|
|
|
# should unlock account and make available for queue
|
|
await pool_mock.unlock(acc.username, Q)
|
|
acc = await pool_mock.get_for_queue(Q)
|
|
assert acc is not None
|
|
assert acc.locks[Q] is not None
|
|
|
|
# should update lock time
|
|
end_time = utc.ts() + 60 # + 1 minute
|
|
await pool_mock.lock_until(acc.username, Q, end_time)
|
|
|
|
acc = await pool_mock.get(acc.username)
|
|
assert int(acc.locks[Q].timestamp()) == end_time
|
|
|
|
|
|
async def test_get_stats(pool_mock: AccountsPool):
|
|
Q = "SearchTimeline"
|
|
|
|
# should return empty stats
|
|
stats = await pool_mock.stats()
|
|
for k, v in stats.items():
|
|
assert v == 0, f"{k} should be 0"
|
|
|
|
# should increate total
|
|
await pool_mock.add_account("user1", "pass1", "email1", "email_pass1")
|
|
stats = await pool_mock.stats()
|
|
assert stats["total"] == 1
|
|
assert stats["active"] == 0
|
|
|
|
# should increate active
|
|
await pool_mock.set_active("user1", True)
|
|
stats = await pool_mock.stats()
|
|
assert stats["total"] == 1
|
|
assert stats["active"] == 1
|
|
|
|
# should update queue stats
|
|
acc = await pool_mock.get_for_queue(Q)
|
|
assert acc is not None
|
|
stats = await pool_mock.stats()
|
|
assert stats["total"] == 1
|
|
assert stats["active"] == 1
|
|
assert stats[f"locked_{Q}"] == 1
|
|
|
|
|
|
async def test_delete_accounts(pool_mock: AccountsPool):
|
|
await pool_mock.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool_mock.add_account("user2", "pass2", "email2", "ep2")
|
|
|
|
await pool_mock.delete_accounts("user1")
|
|
usernames = {x.username for x in await pool_mock.get_all()}
|
|
assert "user1" not in usernames
|
|
assert "user2" in usernames
|
|
|
|
await pool_mock.delete_accounts(["user2"])
|
|
assert len(await pool_mock.get_all()) == 0
|
|
|
|
|
|
async def test_username_lists_are_parameterized(pool_mock: AccountsPool, monkeypatch):
|
|
username = 'user"quoted'
|
|
await pool_mock.add_account(username, "pass", "email", "ep")
|
|
|
|
async def fake_login(account):
|
|
return True
|
|
|
|
monkeypatch.setattr(pool_mock, "login", fake_login)
|
|
assert await pool_mock.login_all([username]) == {"total": 1, "success": 1, "failed": 0}
|
|
|
|
await pool_mock.relogin(username)
|
|
assert (await pool_mock.get(username)).active is False
|
|
|
|
await pool_mock.delete_accounts([username])
|
|
assert await pool_mock.get_all() == []
|
|
|
|
|
|
async def test_login_all_empty_username_list_is_noop(pool_mock: AccountsPool):
|
|
assert await pool_mock.login_all([]) == {"total": 0, "success": 0, "failed": 0}
|
|
|
|
|
|
async def test_relogin_skips_cookie_accounts(pool_mock: AccountsPool, monkeypatch):
|
|
await pool_mock.add_account_cookies("user1", "auth_token=token; ct0=csrf")
|
|
|
|
async def fail_if_called(account):
|
|
pytest.fail("cookie account entered password login")
|
|
|
|
monkeypatch.setattr(pool_mock, "login", fail_if_called)
|
|
await pool_mock.relogin("user1")
|
|
|
|
assert (await pool_mock.get("user1")).active is True
|
|
|
|
|
|
async def test_delete_inactive(pool_mock: AccountsPool):
|
|
await pool_mock.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool_mock.add_account("user2", "pass2", "email2", "ep2")
|
|
await pool_mock.set_active("user2", True)
|
|
|
|
await pool_mock.delete_inactive()
|
|
accs = await pool_mock.get_all()
|
|
assert len(accs) == 1
|
|
assert accs[0].username == "user2"
|
|
|
|
|
|
async def test_reset_locks(pool_mock: AccountsPool):
|
|
Q = "TestQueue"
|
|
await pool_mock.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool_mock.set_active("user1", True)
|
|
|
|
await pool_mock.get_for_queue(Q)
|
|
user = await pool_mock.get("user1")
|
|
assert Q in user.locks
|
|
|
|
await pool_mock.reset_locks()
|
|
user = await pool_mock.get("user1")
|
|
assert Q not in user.locks
|
|
|
|
|
|
async def test_mark_inactive(pool_mock: AccountsPool):
|
|
await pool_mock.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool_mock.set_active("user1", True)
|
|
|
|
await pool_mock.mark_inactive("user1", "banned by system")
|
|
acc = await pool_mock.get("user1")
|
|
assert acc.active is False
|
|
assert acc.error_msg == "banned by system"
|
|
|
|
|
|
async def test_reactivate(pool_mock: AccountsPool):
|
|
await pool_mock.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool_mock.mark_inactive("user1", "banned by system")
|
|
|
|
await pool_mock.reactivate("user1")
|
|
acc = await pool_mock.get("user1")
|
|
assert acc.active is True
|
|
assert acc.error_msg is None
|
|
|
|
# unknown username is a warning, not an error
|
|
await pool_mock.reactivate("nope")
|
|
|
|
|
|
async def test_next_available_at_none_when_empty(pool_mock: AccountsPool):
|
|
assert await pool_mock.next_available_at("TestQueue") is None
|
|
|
|
|
|
async def test_next_available_at_returns_future_time(pool_mock: AccountsPool):
|
|
Q = "TestQueue"
|
|
await pool_mock.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool_mock.set_active("user1", True)
|
|
await pool_mock.lock_until("user1", Q, utc.ts() + 3600)
|
|
|
|
result = await pool_mock.next_available_at(Q)
|
|
assert result is not None
|
|
assert result != "now"
|
|
|
|
|
|
async def test_get_for_queue_or_wait_raises_when_flag_set(pool_mock: AccountsPool):
|
|
pool = AccountsPool(pool_mock._db_file, raise_when_no_account=True)
|
|
with pytest.raises(NoAccountError):
|
|
await pool.get_for_queue_or_wait("TestQueue")
|
|
|
|
|
|
async def test_get_for_queue_or_wait_raises_via_env(pool_mock: AccountsPool, monkeypatch):
|
|
monkeypatch.setenv("TWS_RAISE_WHEN_NO_ACCOUNT", "1")
|
|
with pytest.raises(NoAccountError):
|
|
await pool_mock.get_for_queue_or_wait("TestQueue")
|
|
|
|
|
|
async def test_get_for_queue_or_wait_waits_for_locked_account(pool_mock: AccountsPool, monkeypatch):
|
|
queue = "TestQueue"
|
|
pool = AccountsPool(pool_mock._db_file, wait_timeout=1, wait_interval=0.1)
|
|
await pool.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool.set_active("user1", True)
|
|
await pool.get_for_queue(queue)
|
|
|
|
intervals = []
|
|
|
|
async def release_account(interval):
|
|
intervals.append(interval)
|
|
await pool.unlock("user1", queue)
|
|
|
|
monkeypatch.setattr("twscrape.accounts_pool.asyncio.sleep", release_account)
|
|
|
|
account = await pool.get_for_queue_or_wait(queue)
|
|
|
|
assert account is not None
|
|
assert account.username == "user1"
|
|
assert intervals == [0.1]
|
|
|
|
|
|
async def test_get_for_queue_or_wait_returns_none_when_timeout_is_zero(pool_mock: AccountsPool):
|
|
queue = "TestQueue"
|
|
pool = AccountsPool(pool_mock._db_file, wait_timeout=0)
|
|
await pool.add_account("user1", "pass1", "email1", "ep1")
|
|
await pool.set_active("user1", True)
|
|
await pool.get_for_queue(queue)
|
|
|
|
assert await pool.get_for_queue_or_wait(queue) is None
|
|
|
|
|
|
async def test_get_for_queue_or_wait_does_not_wait_without_active_accounts(
|
|
pool_mock: AccountsPool, monkeypatch
|
|
):
|
|
pool = AccountsPool(pool_mock._db_file, wait_timeout=1)
|
|
|
|
async def fail_if_called(_):
|
|
pytest.fail("should not wait when no accounts are active")
|
|
|
|
monkeypatch.setattr("twscrape.accounts_pool.asyncio.sleep", fail_if_called)
|
|
|
|
assert await pool.get_for_queue_or_wait("TestQueue") is None
|
|
|
|
|
|
def test_api_passes_wait_config_to_new_pool(tmp_path):
|
|
api = API(str(tmp_path / "test.db"), wait_timeout=1, wait_interval=0.1)
|
|
|
|
assert api.pool._wait_timeout == 1
|
|
assert api.pool._wait_interval == 0.1
|
|
|
|
|
|
async def test_accounts_info_active_first(pool_mock: AccountsPool):
|
|
await pool_mock.add_account("user_b", "pass", "email", "ep")
|
|
await pool_mock.add_account("user_a", "pass", "email", "ep")
|
|
await pool_mock.set_active("user_a", True)
|
|
|
|
items = await pool_mock.accounts_info()
|
|
assert items[0]["username"] == "user_a"
|
|
assert items[0]["active"] is True
|
|
assert items[1]["username"] == "user_b"
|
|
assert items[1]["active"] is False
|
|
|
|
|
|
async def test_load_from_file(pool_mock: AccountsPool, tmp_path):
|
|
filepath = tmp_path / "accounts.txt"
|
|
filepath.write_text("user1:pass1:email1:ep1\nuser2:pass2:email2:ep2\n")
|
|
|
|
await pool_mock.load_from_file(str(filepath), "username:password:email:email_password")
|
|
usernames = {x.username for x in await pool_mock.get_all()}
|
|
assert "user1" in usernames
|
|
assert "user2" in usernames
|