o2o-negosium-original/agent/tests/test_h1_qtable_policy.py

"""H1 검증 — 실제 UCB Q-Table 정책 + learning 스키마 영속화 + 온라인 학습.

1. UCB select: 가용 중 UCB 최대 선택, 사용 액션 마스킹, propensity = (1-ε)+ε/n.
2. Q-learning update: Q 가 보상 방향으로 이동, done 시 부트스트랩 없음.
3. predict_action_dist: ε-greedy 분포 합=1, greedy 에 질량.
4. warm_start: 차원 일치 시 Q 복제·visit 감쇠, 불일치 시 ValueError.
5. (DB) 버전 확보 idempotent + 셀 upsert/load 라운드트립.
6. (DB) service.step 반복 호출 시 visit/Q 누적 → 학습 진행 + 테넌트 격리.
"""

import math
import os

import numpy as np
import pytest

from negotiation.policies.base import EpisodeState, PolicyContext, Transition
from negotiation.policies.qtable_policy import UCBQTablePolicy
from negotiation.qtable.domain.model.q_table import QTable
from negotiation.qtable.domain.model.snapshot import NegotiationOutcome, NegotiationSnapshot
from negotiation.qtable.infra.repository.learning_repository import LearningRepository
from services.negotiation_service import NegotiationService
from router.v1.negotiation.protocol import Req_NegotiationStep
from tenancy.config_loader import TenantConfigLoader
from tenancy.registry import TenantEngineRegistry

_TENANTS_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tenants")


def _snap(**o):
    base = dict(revenue_amount=1, distribution_code="A", partner_count=1, acceptance_ratio=0.1,
                input_price=900, anchor_price=800, target_price=1000)
    base.update(o)
    return NegotiationSnapshot(**base)


def _ctx(state_index, n=4, used=None):
    return PolicyContext(state_index=state_index, snapshot=_snap(), action_space_size=n,
                         episode=EpisodeState(used_action_ids=set(used or [])))


# ---- 유닛 ----------------------------------------------------------------
def test_ucb_selects_highest_and_masks_used():
    qt = QTable(2, 4)
    qt.q[0] = np.array([1.0, 9.0, 2.0, 0.0])
    qt.visits[0] = np.array([5, 5, 5, 5])  # 방문 균등 → Q 가 지배
    p = UCBQTablePolicy(qt, exploration_constant=0.1, epsilon=0.1)
    d = p.select(_ctx(0))
    assert d.action_id == 1  # 최고 Q
    assert d.propensity == pytest.approx(0.9 + 0.1 / 4)
    # 1 을 사용처리하면 다음은 다른 액션
    d2 = p.select(_ctx(0, used=[1]))
    assert d2.action_id != 1


def test_ucb_exploration_bonus_prefers_unvisited():
    qt = QTable(1, 3)
    qt.q[0] = np.array([1.0, 0.0, 0.0])
    qt.visits[0] = np.array([100, 0, 0])  # action0 Q 높지만 과방문
    p = UCBQTablePolicy(qt, exploration_constant=math.sqrt(2), epsilon=0.1)
    d = p.select(_ctx(0, n=3))
    assert d.action_id in (1, 2)  # 미방문 액션의 탐색 보너스가 이긴다


def test_q_update_moves_toward_reward():
    qt = QTable(2, 2, learning_rate=0.5)
    new = qt.update(0, 0, reward=10.0, done=True)   # 0 + 0.5*(10-0)=5
    assert new == pytest.approx(5.0)
    # done=False + next state 부트스트랩
    qt.q[1] = np.array([4.0, 0.0])
    new2 = qt.update(0, 1, reward=1.0, next_state_index=1, done=False)  # 0+0.5*(1+0.95*4-0)
    assert new2 == pytest.approx(0.5 * (1 + 0.95 * 4))


def test_predict_action_dist_is_epsilon_greedy():
    qt = QTable(1, 4); qt.q[0] = np.array([0, 9.0, 0, 0]); qt.visits[0] = np.array([5, 5, 5, 5])
    p = UCBQTablePolicy(qt, exploration_constant=0.1, epsilon=0.2)
    dist = p.predict_action_dist(_ctx(0))
    assert dist.sum() == pytest.approx(1.0)
    assert dist[1] == pytest.approx(0.8 + 0.2 / 4)  # greedy


def test_warm_start_dimension_check():
    base = UCBQTablePolicy(QTable(162, 9)); base.qtable.q[0, 0] = 3.0; base.qtable.visits[0, 0] = 10
    tgt = UCBQTablePolicy(QTable(162, 9))
    tgt.warm_start(base)
    assert tgt.qtable.q[0, 0] == 3.0
    assert tgt.qtable.visits[0, 0] == 5  # 감쇠(0.5)
    with pytest.raises(ValueError):
        UCBQTablePolicy(QTable(10, 9)).warm_start(base)  # 차원 불일치


# ---- DB ------------------------------------------------------------------
@pytest.mark.asyncio
async def test_version_and_cell_persistence(db_engine):
    repo = LearningRepository("co-h1")
    v1 = await repo.get_or_create_active_version(state_space_size=162, action_space_size=9, learning_rate=0.1, discount_factor=0.95)
    v2 = await repo.get_or_create_active_version(state_space_size=162, action_space_size=9, learning_rate=0.1, discount_factor=0.95)
    assert v1 == v2  # idempotent (활성 버전 재사용)

    await repo.upsert_cell(v1, state_index=5, action_id=2, q_value=1.5, count=3)
    await repo.upsert_cell(v1, state_index=5, action_id=2, q_value=2.5, count=4)  # 갱신
    qcells, vcells = await repo.load_cells(v1)
    qmap = {(s, a): q for s, a, q in qcells}
    vmap = {(s, a): c for s, a, c in vcells}
    assert qmap[(5, 2)] == 2.5
    assert vmap[(5, 2)] == 4


@pytest.mark.asyncio
async def test_service_step_learns_and_isolates(db_engine):
    reg = TenantEngineRegistry(loader=TenantConfigLoader(tenants_dir=_TENANTS_DIR, cache_ttl_seconds=0))
    eng = await reg.get_engine("imarketkorea")
    svc = NegotiationService()

    def req():
        return Req_NegotiationStep(revenue_amount=20_000_000, distribution_code="A", partner_count=1,
                                   acceptance_ratio=0.11, input_price=990, anchor_price=800, target_price=1000,
                                   round_number=3, outcome="success", log=True, learn=True)

    r1 = await svc.step(eng, req())
    r2 = await svc.step(eng, req())
    r3 = await svc.step(eng, req())
    # UCB 는 미방문 액션을 탐색하므로 매 호출 다른 액션을 고른다(정상). → state 전체 방문이 누적된다.
    assert {r1.action_id, r2.action_id, r3.action_id} == {r1.action_id} or len({r1.action_id, r2.action_id, r3.action_id}) >= 2
    assert r1.learned is True and r1.policy == "qtable_ucb"
    assert r1.updated_q > 0.0  # 성공 보상으로 Q 상승

    # DB 에서 state 전체 방문 누적 확인 (state index 는 imarketkorea config 로 동적 계산)
    from negotiation.qtable.domain.service.state_calculator import state_index
    sidx = state_index(_snap(revenue_amount=20_000_000, acceptance_ratio=0.11, input_price=990, round_number=3), eng.config.state)
    repo_a = LearningRepository("imarketkorea")
    vid = await repo_a.get_or_create_active_version(state_space_size=162, action_space_size=9, learning_rate=0.1, discount_factor=0.95)
    _, vcells = await repo_a.load_cells(vid)
    state_total = sum(c for s, a, c in vcells if s == sidx)
    assert state_total == 3  # 3회 호출 → state 누적 방문 3

    # 테넌트 격리: 자동 온보딩 고객사(UUID)는 별도 학습/별도 state
    other = "00000000-0000-0000-0000-0000000000c1"
    eng2 = await reg.get_engine(other)
    ri = await svc.step(eng2, req())
    assert ri.visit_count == 1

    err, ck = await repo_a.read(lambda s: repo_a.count_experience(s))
    err, ci = await LearningRepository(other).read(lambda s: LearningRepository(other).count_experience(s))
    assert ck == 3 and ci == 1  # experience 격리