"""
Tests for OMOPClient

Tests WebAPI client functionality.
"""

import pytest
from smart_omop import OMOPClient


BASE_URL = "http://ec2-13-49-5-87.eu-north-1.compute.amazonaws.com:8080/WebAPI"


def test_client_initialization():
    """Test client can be initialized."""
    client = OMOPClient(BASE_URL)
    assert client.base_url == BASE_URL
    assert client.timeout == 30
    assert client.max_retries == 3


def test_context_manager():
    """Test client works as context manager."""
    with OMOPClient(BASE_URL) as client:
        assert client is not None


def test_get_sources():
    """Test fetching data sources."""
    with OMOPClient(BASE_URL) as client:
        sources = client.get_sources()
        assert isinstance(sources, list)
        assert len(sources) > 0
        assert 'sourceKey' in sources[0]


def test_get_cohort():
    """Test fetching cohort definition."""
    with OMOPClient(BASE_URL) as client:
        cohort = client.get_cohort(cohort_id=145)
        assert isinstance(cohort, dict)
        assert 'id' in cohort
        assert cohort['id'] == 145


def test_get_generation_status():
    """Test getting generation status."""
    with OMOPClient(BASE_URL) as client:
        statuses = client.get_generation_status(cohort_id=145, source_key="KAGGLECOPD")
        assert isinstance(statuses, list)
        if statuses:
            assert 'status' in statuses[0]


def test_get_cohort_results():
    """Test getting cohort results."""
    with OMOPClient(BASE_URL) as client:
        results = client.get_cohort_results(cohort_id=145, source_key="KAGGLECOPD")
        assert isinstance(results, dict)
        assert 'personCount' in results or 'status' in results
