import pytest
from app.worker import classify_transactions_batch
from google import genai
from unittest.mock import MagicMock, patch

pytestmark = pytest.mark.asyncio


async def test_classify_transactions_batch_retry_success():
    # Mock GenAI client and interaction creation
    mock_client = MagicMock(spec=genai.Client)
    mock_interaction = MagicMock()
    mock_interaction.output_text = '{"classifications": [{"txn_id": "TX1", "category": "Food"}, {"txn_id": "TX2", "category": "Shopping"}]}'
    mock_client.interactions.create.return_value = mock_interaction

    batch_txns = [
        {"txn_id": "TX1", "merchant": "Swiggy", "amount": 150.0},
        {"txn_id": "TX2", "merchant": "Amazon", "amount": 1200.0},
    ]

    results = await classify_transactions_batch(
        mock_client, batch_txns, ["Food", "Shopping", "Other"]
    )
    assert len(results) == 2
    assert results["TX1"]["category"] == "Food"
    assert results["TX1"]["failed"] is False
    assert results["TX2"]["category"] == "Shopping"
    assert results["TX2"]["failed"] is False


async def test_classify_transactions_batch_retry_failure():
    mock_client = MagicMock(spec=genai.Client)
    # Simulate API call failure (e.g. invalid API key or connection error)
    mock_client.interactions.create.side_effect = Exception("API Key Error")

    batch_txns = [
        {"txn_id": "TX1", "merchant": "Swiggy", "amount": 150.0},
    ]

    # Patch sleep to make retry test run instantly
    with patch("asyncio.sleep", return_value=None):
        results = await classify_transactions_batch(
            mock_client, batch_txns, ["Food", "Shopping", "Other"]
        )

    assert len(results) == 1
    assert results["TX1"]["category"] == "Other"
    assert results["TX1"]["failed"] is True
