48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
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
|