from __future__ import annotations

import pytest

import agent_framework.guardrails.framework_llm_client as module


class _OwnedLLM:
    def __init__(self):
        self.closed = False

    async def ainvoke(self, *args, **kwargs):
        return '{"allowed": true, "reason": ""}'

    async def aclose(self):
        self.closed = True


@pytest.mark.asyncio
async def test_temporary_guardrail_provider_is_closed_inside_owning_event_loop(monkeypatch):
    llm = _OwnedLLM()
    monkeypatch.setattr(module, '_ensure_framework_llm', lambda value: llm)
    monkeypatch.setenv('USE_MOCK_LLM', 'false')

    result = await module.classify_with_framework_llm(
        None,
        'AOFERTA',
        {'text': 'Posso ajudar?', 'context': {}},
    )

    assert result['allowed'] is True
    assert llm.closed is True


@pytest.mark.asyncio
async def test_caller_owned_guardrail_provider_is_not_closed(monkeypatch):
    llm = _OwnedLLM()
    monkeypatch.setenv('USE_MOCK_LLM', 'false')

    result = await module.classify_with_framework_llm(
        llm,
        'AOFERTA',
        {'text': 'Posso ajudar?', 'context': {}},
    )

    assert result['allowed'] is True
    assert llm.closed is False