agent-framework/python/packages/azurefunctions/tests/test_entities.py

# Copyright (c) Microsoft. All rights reserved.

"""Unit tests for create_agent_entity factory function.

Run with: pytest tests/test_entities.py -v
"""

from collections.abc import Callable
from typing import Any, TypeVar
from unittest.mock import AsyncMock, Mock

import pytest
from agent_framework import AgentResponse, Message

from agent_framework_azurefunctions._entities import create_agent_entity

FuncT = TypeVar("FuncT", bound=Callable[..., Any])


def _agent_response(text: str | None) -> AgentResponse:
    """Create an AgentResponse with a single assistant message."""
    message = (
        Message(role="assistant", contents=[text]) if text is not None else Message(role="assistant", contents=[""])
    )
    return AgentResponse(messages=[message])


class TestCreateAgentEntity:
    """Test suite for the create_agent_entity factory function."""

    def test_create_agent_entity_returns_callable(self) -> None:
        """Test that create_agent_entity returns a callable."""
        mock_agent = Mock()

        entity_function = create_agent_entity(mock_agent)

        assert callable(entity_function)

    def test_entity_function_handles_run_agent(self) -> None:
        """Test that the entity function handles the run_agent operation."""
        mock_agent = Mock()
        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))

        entity_function = create_agent_entity(mock_agent)

        # Mock context
        mock_context = Mock()
        mock_context.operation_name = "run"
        mock_context.entity_key = "conv-123"
        mock_context.get_input.return_value = {
            "message": "Test message",
            "correlationId": "corr-entity-factory",
        }
        mock_context.get_state.return_value = None

        # Execute
        entity_function(mock_context)

        # Verify result and state were set
        assert mock_context.set_result.called
        assert mock_context.set_state.called

    def test_entity_function_handles_reset(self) -> None:
        """Test that the entity function handles the reset operation."""
        mock_agent = Mock()

        entity_function = create_agent_entity(mock_agent)

        # Mock context with existing state
        mock_context = Mock()
        mock_context.operation_name = "reset"
        mock_context.get_state.return_value = {
            "schemaVersion": "1.0.0",
            "data": {
                "conversationHistory": [
                    {
                        "$type": "request",
                        "correlationId": "test-correlation-id",
                        "createdAt": "2024-01-01T00:00:00Z",
                        "messages": [
                            {
                                "role": "user",
                                "contents": [{"$type": "text", "text": "test"}],
                            }
                        ],
                    }
                ]
            },
        }

        # Execute
        entity_function(mock_context)

        # Verify reset result
        assert mock_context.set_result.called
        result = mock_context.set_result.call_args[0][0]
        assert result["status"] == "reset"

        # Verify state was cleared
        assert mock_context.set_state.called
        state = mock_context.set_state.call_args[0][0]
        assert state["data"]["conversationHistory"] == []

    def test_entity_function_handles_unknown_operation(self) -> None:
        """Test that the entity function handles unknown operations."""
        mock_agent = Mock()

        entity_function = create_agent_entity(mock_agent)

        mock_context = Mock()
        mock_context.operation_name = "invalid_operation"
        mock_context.get_state.return_value = None

        # Execute
        entity_function(mock_context)

        # Verify error result
        assert mock_context.set_result.called
        result = mock_context.set_result.call_args[0][0]
        assert "error" in result
        assert "invalid_operation" in result["error"].lower()

    def test_entity_function_creates_new_entity_on_first_call(self) -> None:
        """Test that the entity function creates a new entity when no state exists."""
        mock_agent = Mock()
        mock_agent.__class__.__name__ = "Agent"

        entity_function = create_agent_entity(mock_agent)
        mock_context = Mock()
        mock_context.operation_name = "reset"
        mock_context.get_state.return_value = None  # No existing state

        # Execute
        entity_function(mock_context)

        # Verify new entity state was created
        assert mock_context.set_result.called
        result = mock_context.set_result.call_args[0][0]
        assert result["status"] == "reset"
        assert mock_context.set_state.called
        state = mock_context.set_state.call_args[0][0]
        assert state["data"] == {"conversationHistory": []}

    def test_entity_function_restores_existing_state(self) -> None:
        """Test that the entity function can operate when existing state is present."""
        mock_agent = Mock()

        entity_function = create_agent_entity(mock_agent)

        existing_state = {
            "schemaVersion": "1.0.0",
            "data": {
                "conversationHistory": [
                    {
                        "$type": "request",
                        "correlationId": "corr-existing-1",
                        "createdAt": "2024-01-01T00:00:00Z",
                        "messages": [
                            {
                                "role": "user",
                                "contents": [
                                    {
                                        "$type": "text",
                                        "text": "msg1",
                                    }
                                ],
                            }
                        ],
                    },
                    {
                        "$type": "response",
                        "correlationId": "corr-existing-1",
                        "createdAt": "2024-01-01T00:05:00Z",
                        "messages": [
                            {
                                "role": "assistant",
                                "contents": [
                                    {
                                        "$type": "text",
                                        "text": "resp1",
                                    }
                                ],
                            }
                        ],
                    },
                ],
            },
        }

        mock_context = Mock()
        mock_context.operation_name = "reset"
        mock_context.get_state.return_value = existing_state

        entity_function(mock_context)

        assert mock_context.set_result.called

        # Reset should clear history and persist via set_state
        assert mock_context.set_state.called
        persisted_state = mock_context.set_state.call_args[0][0]
        assert persisted_state["data"]["conversationHistory"] == []

    def test_entity_function_handles_string_input(self) -> None:
        """Test that the entity function handles non-dict input by converting to string."""
        mock_agent = Mock()
        mock_agent.run = AsyncMock(return_value=_agent_response("String response"))

        entity_function = create_agent_entity(mock_agent)

        # Mock context with non-dict input (like a number)
        mock_context = Mock()
        mock_context.operation_name = "run"
        mock_context.entity_key = "conv-456"
        # Use a number to test the str() conversion path
        mock_context.get_input.return_value = 12345
        mock_context.get_state.return_value = None

        # Execute - entity will convert non-dict input to string
        entity_function(mock_context)

        # Verify the result was set
        assert mock_context.set_result.called

    def test_entity_function_handles_none_input(self) -> None:
        """Test that the entity function handles None input by converting to empty string."""
        mock_agent = Mock()
        mock_agent.run = AsyncMock(return_value=_agent_response("Empty response"))

        entity_function = create_agent_entity(mock_agent)

        # Mock context with None input
        mock_context = Mock()
        mock_context.operation_name = "run"
        mock_context.entity_key = "conv-789"
        mock_context.get_input.return_value = None
        mock_context.get_state.return_value = None

        # Execute - should hit error path since entity expects dict or valid JSON string
        entity_function(mock_context)

        # Verify the result was set (likely error result)
        assert mock_context.set_result.called

    def test_entity_function_handles_event_loop_runtime_error(self) -> None:
        """Test that the entity function handles RuntimeError from get_event_loop by creating a new loop."""
        from unittest.mock import patch

        mock_agent = Mock()
        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))

        entity_function = create_agent_entity(mock_agent)

        mock_context = Mock()
        mock_context.operation_name = "run"
        mock_context.entity_key = "conv-loop-test"
        mock_context.get_input.return_value = {"message": "Test"}
        mock_context.get_state.return_value = None

        # Simulate RuntimeError when getting event loop
        with (
            patch("asyncio.get_event_loop", side_effect=RuntimeError("No event loop")),
            patch("asyncio.new_event_loop") as mock_new_loop,
            patch("asyncio.set_event_loop") as mock_set_loop,
        ):
            mock_loop = Mock()
            mock_loop.is_running.return_value = False
            mock_loop.run_until_complete = Mock()
            mock_new_loop.return_value = mock_loop

            # Execute
            entity_function(mock_context)

            # Verify new event loop was created
            mock_new_loop.assert_called_once()
            mock_set_loop.assert_called_once_with(mock_loop)

    def test_entity_function_handles_running_event_loop(self) -> None:
        """Test that the entity function handles a running event loop by creating a temporary loop."""
        from unittest.mock import patch

        mock_agent = Mock()
        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))

        entity_function = create_agent_entity(mock_agent)

        mock_context = Mock()
        mock_context.operation_name = "run"
        mock_context.entity_key = "conv-running-loop"
        mock_context.get_input.return_value = {"message": "Test"}
        mock_context.get_state.return_value = None

        # Simulate a running event loop
        mock_existing_loop = Mock()
        mock_existing_loop.is_running.return_value = True

        mock_temp_loop = Mock()
        mock_temp_loop.run_until_complete = Mock()
        mock_temp_loop.close = Mock()

        with (
            patch("asyncio.get_event_loop", return_value=mock_existing_loop),
            patch("asyncio.new_event_loop", return_value=mock_temp_loop),
        ):
            # Execute
            entity_function(mock_context)

            # Verify temporary loop was created and closed
            mock_temp_loop.run_until_complete.assert_called_once()
            mock_temp_loop.close.assert_called_once()


if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])