import pytest
from src.vendor_capabilities import VendorCapabilities, get_capabilities, register

@pytest.fixture(autouse=True)
def _clean_registry():
 import src.vendor_capabilities
 snapshot = src.vendor_capabilities._REGISTRY.copy()
 yield
 src.vendor_capabilities._REGISTRY.clear()
 src.vendor_capabilities._REGISTRY.update(snapshot)

def test_registry_lookup_known_model():
 caps = VendorCapabilities(
  vendor='qwen',
  model='qwen-max',
  vision=False,
  context_window=32768
 )
 register(caps)
 retrieved = get_capabilities('qwen', 'qwen-max')
 assert retrieved.vendor == 'qwen'
 assert retrieved.model == 'qwen-max'
 assert retrieved.context_window == 32768
 assert retrieved.vision is False

def test_fallback_to_vendor_default():
 caps = VendorCapabilities(
  vendor='llama',
  model='*',
  context_window=131072,
  cost_tracking=False
 )
 register(caps)
 retrieved = get_capabilities('llama', 'llama-3.3-future-unregistered')
 assert retrieved.context_window == 131072
 assert retrieved.cost_tracking is False

def test_unknown_vendor_raises():
 with pytest.raises(KeyError, match='No capabilities registered'):
  get_capabilities('nonexistent_vendor', 'anymodel')