From 3206d38304c8db9c63aeeb54824e0d0ab63c148f Mon Sep 17 00:00:00 2001 From: agatha Date: Sat, 14 Mar 2026 16:44:16 -0400 Subject: [PATCH] test: add integration tests for acquire and release flow --- tests/integration/test_acquire.py | 202 ++++++++++++++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100644 tests/integration/test_acquire.py diff --git a/tests/integration/test_acquire.py b/tests/integration/test_acquire.py new file mode 100644 index 0000000..2292344 --- /dev/null +++ b/tests/integration/test_acquire.py @@ -0,0 +1,202 @@ +import random +import uuid + +import pytest +from httpx import ASGITransport, AsyncClient + +from proxy_pool.app import create_app +from proxy_pool.proxy.models import Proxy, ProxyProtocol, ProxySource, ProxyStatus + + +@pytest.fixture +async def client(): + app = create_app() + async with app.router.lifespan_context(app): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as client: + yield client + + +async def register_user(client): + email = f"test-{uuid.uuid4().hex[:8]}@example.com" + resp = await client.post( + "/auth/register", + json={"email": email, "display_name": "Test"}, + ) + data = resp.json() + return data["api_key"]["key"] + + +async def seed_proxies(client, count=3): + """Insert active proxies directly via the app's session factory.""" + app = client._transport.app + session_factory = app.state.session_factory + + async with session_factory() as db: + source = ProxySource( + url=f"https://example.com/{uuid.uuid4().hex[:8]}.txt", + parser_name="plaintext", + ) + db.add(source) + await db.flush() + + base = random.randint(1, 250) + proxy_ids = [] + for i in range(count): + proxy = Proxy( + ip=f"198.51.{base}.{i + 1}", + port=9000 + i, + protocol=ProxyProtocol.HTTP, + source_id=source.id, + status=ProxyStatus.ACTIVE, + score=0.9 - (i * 0.1), + country="US", + ) + db.add(proxy) + await db.flush() + proxy_ids.append(str(proxy.id)) + + await db.commit() + return proxy_ids + + +def auth(key: str) -> dict: + return {"Authorization": f"Bearer {key}"} + + +class TestAcquireProxy: + async def test_acquire_returns_proxy_and_debits_credit(self, client): + api_key = await register_user(client) + await seed_proxies(client) + + resp = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={"protocol": "http"}, + ) + + assert resp.status_code == 200 + data = resp.json() + assert data["lease_id"] is not None + assert data["proxy"]["status"] == "active" + assert data["credits_remaining"] == 99 + assert data["expires_at"] is not None + + async def test_acquire_no_matching_proxies(self, client): + api_key = await register_user(client) + + resp = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={"protocol": "socks5", "country": "ZZ"}, + ) + + assert resp.status_code == 404 + + async def test_acquire_unauthenticated(self, client): + resp = await client.post( + "/proxies/acquire", + json={"protocol": "http"}, + ) + + assert resp.status_code == 401 + + async def test_acquire_selects_highest_score(self, client): + api_key = await register_user(client) + await seed_proxies(client, count=3) + + resp = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={}, + ) + + assert resp.status_code == 200 + # First proxy should be the highest scored + assert resp.json()["proxy"]["score"] >= 0.7 + + async def test_acquire_skips_leased_proxies(self, client): + api_key = await register_user(client) + await seed_proxies(client, count=2) + + # Acquire first proxy + resp1 = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={}, + ) + assert resp1.status_code == 200 + first_proxy_id = resp1.json()["proxy"]["id"] + + # Acquire second — should get a different proxy + resp2 = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={}, + ) + assert resp2.status_code == 200 + second_proxy_id = resp2.json()["proxy"]["id"] + + assert first_proxy_id != second_proxy_id + + +class TestReleaseProxy: + async def test_release_makes_proxy_available(self, client): + api_key = await register_user(client) + await seed_proxies(client, count=1) + + # Acquire + acq_resp = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={}, + ) + lease_id = acq_resp.json()["lease_id"] + + # Release + rel_resp = await client.post( + f"/proxies/acquire/{lease_id}/release", + headers=auth(api_key), + ) + + assert rel_resp.status_code == 200 + assert rel_resp.json()["released"] is True + + async def test_release_nonexistent_lease(self, client): + api_key = await register_user(client) + + resp = await client.post( + f"/proxies/acquire/{uuid.uuid4()}/release", + headers=auth(api_key), + ) + + assert resp.status_code == 404 + + async def test_release_then_reacquire(self, client): + api_key = await register_user(client) + await seed_proxies(client, count=1) + + # Acquire + acq1 = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={}, + ) + lease_id = acq1.json()["lease_id"] + proxy_id = acq1.json()["proxy"]["id"] + + # Release + await client.post( + f"/proxies/acquire/{lease_id}/release", + headers=auth(api_key), + ) + + # Acquire again — should get the same proxy since it's the only one + acq2 = await client.post( + "/proxies/acquire", + headers=auth(api_key), + json={}, + ) + + assert acq2.status_code == 200 + assert acq2.json()["proxy"]["id"] == proxy_id \ No newline at end of file