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

"""P5 검증 — 베이스 warm-start / cold-start 3단 (계획서 D).

1. warm_start_from_base: 차원 호환 시 base Q값/방문수 복제(visit 감쇠), base_version_id 추적.
2. cold-start: 신규 테넌트 첫 정책 로드 → v000_warmstart_from_base 생성.
3. 차원 불일치 → warm-start None → 휴리스틱 빈 버전 폴백.
4. 활성 버전 있으면 warm-start 안 함(기존 사용).
"""

import os

import pytest

from common.database.model.models import BASE_COMPANY_ID
from negotiation.policy.model_store import QTablePolicyStore
from negotiation.qtable.infra.repository.learning_repository import LearningRepository
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 _reg():
    return TenantEngineRegistry(loader=TenantConfigLoader(tenants_dir=_TENANTS_DIR, cache_ttl_seconds=0))


async def _seed_base(S=162, A=9):
    """_base 활성 버전 + 셀 시드 (state 5 의 action 2·3)."""
    base = LearningRepository(BASE_COMPANY_ID)
    vid = await base.get_or_create_active_version(
        state_space_size=S, action_space_size=A, learning_rate=0.1, discount_factor=0.95,
        scope=1, version_name="base_v000")
    await base.reset_learning()
    await base.upsert_cell(vid, state_index=5, action_id=2, q_value=0.9, count=10)
    await base.upsert_cell(vid, state_index=5, action_id=3, q_value=0.5, count=4)
    return vid


@pytest.mark.asyncio
async def test_warm_start_copies_base_with_decayed_visits(db_engine):
    base_vid = await _seed_base()
    repo = LearningRepository("co-new")
    vid = await repo.warm_start_from_base(state_space_size=162, action_space_size=9,
                                          learning_rate=0.1, discount_factor=0.95, visit_decay=0.5)
    assert vid is not None

    # 버전 메타: scope=tenant, base_version_id 추적, 활성
    err, ver = await repo.read(lambda s: repo.get_version_by_name(s, "v000_warmstart_from_base"))
    assert ver is not None and ver.scope == 2 and ver.is_active
    assert str(ver.base_version_id) == str(base_vid)

    # Q값 복제 + visit 감쇠 복제
    qcells, vcells = await repo.load_cells(vid)
    qmap = {(s, a): q for s, a, q in qcells}
    vmap = {(s, a): c for s, a, c in vcells}
    assert qmap[(5, 2)] == 0.9 and qmap[(5, 3)] == 0.5
    assert vmap[(5, 2)] == 5 and vmap[(5, 3)] == 2   # 10*0.5, 4*0.5


@pytest.mark.asyncio
async def test_cold_start_creates_warmstart_version(db_engine):
    await _seed_base(A=11)                          # imarketkorea action_space=11 과 차원 일치해야 warm-start
    eng = await _reg().get_engine("imarketkorea")   # 활성 버전 없음 → cold-start
    policy, version_id, repo = await QTablePolicyStore.load(eng)
    err, ver = await repo.read(lambda s: repo.get_active_version(s))
    assert ver.version_name == "v000_warmstart_from_base"
    # 복제된 Q가 정책에 적재됨
    assert policy.qtable.q[5, 2] == 0.9


@pytest.mark.asyncio
async def test_catalog_dim_change_migrates_preserving_learning(db_engine):
    """카탈로그 카드 수 변경(9→11) 시 학습 보존 마이그레이션 — 겹치는 셀 복사 + 새 카드 fresh."""
    import uuid as _uuid
    cid = str(_uuid.uuid4())
    # 이 회사 활성 버전을 A=9 로 시드 + 셀 (5,2)=0.9
    repo = LearningRepository(cid)
    vid = await repo.get_or_create_active_version(
        state_space_size=162, action_space_size=9, learning_rate=0.1, discount_factor=0.95,
        scope=2, version_name="old_v9")
    await repo.upsert_cell(vid, state_index=5, action_id=2, q_value=0.9, count=7)

    # 엔진(_base type:db → 카탈로그 11장) 로드 → 9≠11 감지 → 마이그레이션
    eng = await _reg().get_engine(cid)
    assert eng.action_space_size == 11
    policy, new_vid, _ = await QTablePolicyStore.load(eng)
    assert str(new_vid) != str(vid)            # 새 버전
    assert policy.qtable.q[5, 2] == 0.9        # 기존 학습 보존
    assert policy.qtable.q[5, 10] == 0.0       # 새 카드(action 10) fresh
    assert policy.qtable.visits[5, 2] == 7     # 방문수도 보존
    # 새 버전이 활성 · 차원 11
    err, active = await repo.read(lambda s: repo.get_active_version(s))
    assert str(active.version_id) == str(new_vid) and active.action_space_size == 11


@pytest.mark.asyncio
async def test_migration_remaps_by_card_number(db_engine):
    """카드번호 기반 마이그레이션 — 중간 카드 삭제로 action_id 가 밀려도 학습이 카드를 따라간다."""
    import uuid as _uuid
    cid = str(_uuid.uuid4())
    repo = LearningRepository(cid)
    # A=3, action_cards=[NGC-001, NGC-002, NGC-003]. (5,2)=0.9 는 NGC-003 의 학습.
    vid = await repo.get_or_create_active_version(
        state_space_size=162, action_space_size=3, learning_rate=0.1, discount_factor=0.95,
        scope=2, version_name="v_cards3", action_cards=["NGC-001", "NGC-002", "NGC-003"])
    await repo.upsert_cell(vid, state_index=5, action_id=2, q_value=0.9, count=4)  # NGC-003
    await repo.upsert_cell(vid, state_index=5, action_id=0, q_value=0.3, count=2)  # NGC-001
    _, active = await repo.read(lambda s: repo.get_active_version(s))

    # 새 카탈로그: 중간 NGC-002 제거 → [NGC-001, NGC-003] (A=2). NGC-003: old action 2 → new action 1.
    new_vid = await repo.migrate_active_version_dim(
        old_version=active, state_space_size=162, action_space_size=2,
        learning_rate=0.1, discount_factor=0.95, action_cards=["NGC-001", "NGC-003"])
    assert new_vid is not None
    qcells, _ = await repo.load_cells(new_vid)
    qmap = {(st, a): q for st, a, q in qcells}
    assert qmap.get((5, 1)) == 0.9   # NGC-003 학습이 새 action_id 1 로 따라감(밀림 보정)
    assert qmap.get((5, 0)) == 0.3   # NGC-001 은 그대로 action_id 0
    assert (5, 2) not in qmap        # 삭제된 NGC-002 자리 없음


@pytest.mark.asyncio
async def test_dimension_mismatch_falls_back_to_heuristic(db_engine):
    await _seed_base(S=162, A=9)
    repo = LearningRepository("co-mismatch")
    # 다른 차원 요청 → 복제 불가
    vid = await repo.warm_start_from_base(state_space_size=100, action_space_size=9,
                                          learning_rate=0.1, discount_factor=0.95)
    assert vid is None  # 호출자가 휴리스틱 폴백


@pytest.mark.asyncio
async def test_no_base_returns_none(db_engine):
    # base 미시드 → warm-start None
    repo = LearningRepository("co-nobase")
    vid = await repo.warm_start_from_base(state_space_size=162, action_space_size=9,
                                          learning_rate=0.1, discount_factor=0.95)
    assert vid is None


@pytest.mark.asyncio
async def test_existing_version_not_warmstarted(db_engine):
    await _seed_base()
    eng = await _reg().get_engine("imarketkorea")
    # 첫 로드 → warm-start 버전 생성
    await QTablePolicyStore.load(eng)
    # 둘째 로드 → 기존 활성 버전 재사용(중복 warm-start 안 함)
    _, _, repo = await QTablePolicyStore.load(eng)
    err, vers = await repo.read(lambda s: repo.list_versions(s))
    assert len(vers) == 1