o2o-negosium-original/lps/tests/test_port_lease.py

"""프로세스 간 공유 포트 장부(proxy_port) 테스트 — 실제 DB 필요.

여기서 지키는 계약은 멀티 프로세스 운영과 직결된다:
  - 두 프로세스가 같은 IP 를 동시에 잡으면 그 IP 가 두 배 속도로 탄다
  - 한 프로세스가 태운 IP 를 다른 프로세스가 곧바로 집으면 차단이 전파된다
  - 워커가 죽은 채로 임대를 쥐고 있으면 그 IP 는 영영 안 돌아온다(임대 만료로 자동 회수돼야 함)
"""

import pytest
from sqlalchemy import text

from common.database.db_session_manager import DB_SESSION_MNG
from common.enums import DBType, DBWRType
from crud.port_lease import PortLeaseStore

HOST_A, HOST_B = "test-gate.example", "test-kr.example"
P_START, P_END = 20001, 20005


@pytest.fixture
async def store(db_engine):
    st = PortLeaseStore()
    await _exec("DELETE FROM proxy_port WHERE host IN (:a, :b)", {"a": HOST_A, "b": HOST_B})
    await st.ensure_ports(HOST_A, P_START, P_END)
    await st.ensure_ports(HOST_B, P_START, P_END)
    yield st
    await _exec("DELETE FROM proxy_port WHERE host IN (:a, :b)", {"a": HOST_A, "b": HOST_B})


async def _exec(sql: str, params: dict | None = None):
    s = await DB_SESSION_MNG.start_session(DBType.MAIN.value, DBWRType.DB_WRITE.value)
    try:
        await s.execute(text(sql), params or {})
        await s.commit()
    finally:
        await DB_SESSION_MNG.end_session(DBType.MAIN.value, DBWRType.DB_WRITE.value)


# ── 종료 시 반납 (2026-08-05) ────────────────────────────────────────────
async def test_proxy_release_returns_the_port_immediately(store):
    """워커 종료가 임대를 놓아야 한다 — 안 놓으면 재시작해도 최대 sticky 수명(10분)까지 못 쓴다."""
    from services.search.proxy import DecodoProxy

    px = DecodoProxy(host=HOST_A, store=store, owner="shutdown-w0")
    px.username, px.password = "u", "p"
    px.port_start, px.port_end = P_START, P_END
    port = await px.ensure_port()
    assert port is not None
    assert (await store.snapshot(HOST_A))[HOST_A]["held"] == 1

    await px.release()
    assert (await store.snapshot(HOST_A))[HOST_A]["held"] == 0, "종료 후에는 남이 곧바로 쓸 수 있어야 한다"
    assert px._leased is None
    assert await store.acquire(HOST_A, "other-process", 600) is not None


async def test_proxy_release_is_safe_without_a_lease(store):
    """임대를 못 잡은 채 종료돼도(풀 고갈 등) 예외 없이 지나가야 한다."""
    from services.search.proxy import DecodoProxy

    px = DecodoProxy(host=HOST_A, store=store, owner="never-leased")
    px.username, px.password = "u", "p"
    px.port_start, px.port_end = P_START, P_END
    await px.release()


# ── 배타 임대(프로세스 간) ───────────────────────────────────────────────
async def test_no_two_owners_get_the_same_port(store):
    """서로 다른 프로세스를 흉내낸 owner 5개가 5포트를 하나씩 나눠 가져야 한다."""
    got = [await store.acquire(HOST_A, f"p{i}-w0", 600) for i in range(5)]
    assert None not in got
    assert len(set(got)) == 5


async def test_pool_exhaustion_returns_none(store):
    for i in range(5):
        assert await store.acquire(HOST_A, f"p{i}", 600) is not None
    assert await store.acquire(HOST_A, "p-late", 600) is None   # 남는 포트 없음 → 호출부가 대기/재시도


async def test_gateways_are_independent(store):
    """같은 포트 번호라도 게이트웨이가 다르면 다른 IP(실측) — 한쪽 임대가 다른 쪽을 막으면 안 된다."""
    a = await store.acquire(HOST_A, "coupang", 600)
    b = await store.acquire(HOST_B, "naver", 600)
    assert a == b == P_START          # 번호는 같아도 서로 다른 자원
    assert (await store.available(HOST_A)) == 4 and (await store.available(HOST_B)) == 4


# ── 차단·휴식(전역) ──────────────────────────────────────────────────────
async def test_burned_port_is_invisible_to_every_process(store):
    p = await store.acquire(HOST_A, "coupang-p1", 600)
    await store.burn(HOST_A, p, 1800, owner="coupang-p1")
    others = {await store.acquire(HOST_A, f"naver-p{i}", 600) for i in range(4)}
    assert p not in others            # 다른 프로세스도 태운 IP 를 못 집는다
    assert await store.acquire(HOST_A, "naver-p9", 600) is None


async def test_rested_port_is_held_back_then_returns(store):
    p = await store.acquire(HOST_A, "naver-p1", 600)
    await store.release(HOST_A, p, "naver-p1", rest_sec=600)
    others = {await store.acquire(HOST_A, f"x{i}", 600) for i in range(4)}
    assert p not in others            # 쉬는 동안은 아무도 못 집는다
    await store.release(HOST_A, p, "nobody", rest_sec=0)   # 소유자 불일치 → 무시돼야 함
    assert await store.acquire(HOST_A, "y", 600) is None
    await _exec("UPDATE proxy_port SET rest_until = now() - interval '1 s' WHERE host=:h AND port=:p",
                {"h": HOST_A, "p": p})
    assert await store.acquire(HOST_A, "z", 600) == p      # 휴식 만료 → 복귀


async def test_release_without_rest_returns_port_to_pool(store):
    """반납된 포트는 곧바로 풀에 돌아온다. 단 **다음 차례가 되는 건 아니다** —
    LRU 라 한 번도 안 쓴 포트가 먼저 나가고, 반납분은 한 바퀴 뒤에 다시 온다(회전의 정의)."""
    p = await store.acquire(HOST_A, "naver-p1", 600)
    await store.release(HOST_A, p, "naver-p1")
    assert await store.available(HOST_A) == 5                  # 5개 전부 가용

    picked = [await store.acquire(HOST_A, f"w{i}", 600) for i in range(5)]
    assert p in picked and picked[-1] == p                     # 가장 최근 사용분이 맨 마지막


# ── 죽은 프로세스 회수 ───────────────────────────────────────────────────
async def test_expired_lease_is_reclaimed_without_a_reaper(store):
    """워커가 임대를 쥔 채 죽어도 sticky 수명이 지나면 자동으로 풀려야 한다."""
    p = await store.acquire(HOST_A, "dead-process", 600)
    assert await store.available(HOST_A) == 4                  # 죽은 프로세스가 쥔 동안은 빠져 있고
    await _exec("UPDATE proxy_port SET leased_until = now() - interval '1 s' WHERE host=:h AND port=:p",
                {"h": HOST_A, "p": p})
    assert await store.available(HOST_A) == 5                  # 임대 만료 → 별도 정리 없이 자동 복귀
    picked = [await store.acquire(HOST_A, f"alive{i}", 600) for i in range(5)]
    assert p in picked


async def test_renew_keeps_the_lease_and_fails_after_takeover(store):
    p = await store.acquire(HOST_A, "p1", 600)
    assert await store.renew(HOST_A, p, "p1", 600) is True
    await _exec("UPDATE proxy_port SET owner='p2' WHERE host=:h AND port=:p", {"h": HOST_A, "p": p})
    assert await store.renew(HOST_A, p, "p1", 600) is False   # 남이 가져감 → 새로 잡아야 함


# ── 회전(LRU) ────────────────────────────────────────────────────────────
async def test_rotation_is_least_recently_used(store):
    """프로세스가 몇 개든 '가장 오래 안 쓴 IP'를 집으므로 전체가 한 바퀴씩 돈다."""
    first = []
    for i in range(5):                      # 5포트를 한 바퀴 소진
        p = await store.acquire(HOST_A, f"w{i}", 600)
        first.append(p)
        await store.release(HOST_A, p, f"w{i}")
    second = []
    for i in range(5):
        p = await store.acquire(HOST_A, f"w{i}", 600)
        second.append(p)
        await store.release(HOST_A, p, f"w{i}")
    assert sorted(first) == sorted(second)  # 같은 5개를
    assert first == second                  # 같은 순서로(LRU) — 특정 IP 편중 없음


async def test_snapshot_counts_each_state(store):
    held = await store.acquire(HOST_A, "w0", 600)
    rested = await store.acquire(HOST_A, "w1", 600)
    await store.release(HOST_A, rested, "w1", rest_sec=600)
    burned = await store.acquire(HOST_A, "w2", 600)
    await store.burn(HOST_A, burned, 1800)
    snap = (await store.snapshot(HOST_A))[HOST_A]
    assert snap["held"] == 1 and snap["resting"] == 1 and snap["cooling"] == 1
    assert snap["available"] == 2 and snap["total"] == 5
    assert held != rested != burned