@cdklabs/cdk-appmod-catalog-blueprints 1.10.0 → 1.12.0

package/lib/framework/agents/resources/interactive-agent-handler/test_handler.py

@@ -2,22 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Unit tests for Interactive Agent Lambda Handler
+Unit tests for Interactive Agent FastAPI Handler with Strands-native session/context.
 """
 
 import pytest
 import json
 from unittest.mock import Mock, patch, AsyncMock, MagicMock
-from index import (
-    handler,
-    handle_connect,
-    handle_disconnect,
-    handle_message,
-    SessionManager,
-    ContextManager,
-    send_to_connection,
-    load_system_prompt
-)
 
 
 # Fixtures
@@ -29,71 +19,6 @@ def mock_s3_client():
         yield mock
 
 
-@pytest.fixture
-def mock_apigateway_client():
-    """Mock API Gateway Management API client"""
-    with patch('index.apigateway_client') as mock:
-        yield mock
-
-
-@pytest.fixture
-def connect_event():
-    """Sample $connect event"""
-    return {
-        'requestContext': {
-            'routeKey': '$connect',
-            'connectionId': 'test-connection-123',
-            'requestTimeEpoch': 1234567890,
-            'authorizer': {
-                'principalId': 'user-123'
-            }
-        }
-    }
-
-
-@pytest.fixture
-def disconnect_event():
-    """Sample $disconnect event"""
-    return {
-        'requestContext': {
-            'routeKey': '$disconnect',
-            'connectionId': 'test-connection-123',
-            'authorizer': {
-                'principalId': 'user-123'
-            }
-        }
-    }
-
-
-@pytest.fixture
-def message_event():
-    """Sample message event"""
-    return {
-        'requestContext': {
-            'routeKey': '$default',
-            'connectionId': 'test-connection-123',
-            'domainName': 'test.execute-api.us-east-1.amazonaws.com',
-            'stage': 'prod',
-            'authorizer': {
-                'principalId': 'user-123'
-            }
-        },
-        'body': json.dumps({
-            'message': 'Hello, agent!'
-        })
-    }
-
-
-@pytest.fixture
-def lambda_context():
-    """Mock Lambda context"""
-    context = Mock()
-    context.function_name = 'test-function'
-    context.memory_limit_in_mb = 1024
-    context.invoked_function_arn = 'arn:aws:lambda:us-east-1:123456789012:function:test-function'
-    return context
-
-
 # Tests for load_system_prompt
 
 def test_load_system_prompt_success(mock_s3_client):
@@ -101,13 +26,14 @@ def test_load_system_prompt_success(mock_s3_client):
     mock_s3_client.get_object.return_value = {
         'Body': Mock(read=Mock(return_value=b'You are a helpful assistant.'))
     }
-    
+
     with patch.dict('os.environ', {
         'SYSTEM_PROMPT_S3_BUCKET_NAME': 'test-bucket',
         'SYSTEM_PROMPT_S3_KEY': 'prompts/system.txt'
     }):
+        from index import load_system_prompt
         prompt = load_system_prompt()
-    
+
     assert prompt == 'You are a helpful assistant.'
     mock_s3_client.get_object.assert_called_once_with(
         Bucket='test-bucket',
@@ -118,384 +44,259 @@ def test_load_system_prompt_success(mock_s3_client):
 def test_load_system_prompt_missing_config():
     """Test system prompt loading with missing configuration"""
     with patch.dict('os.environ', {}, clear=True):
+        from index import load_system_prompt
         prompt = load_system_prompt()
-    
+
     assert prompt == "You are a helpful AI assistant."
 
 
 def test_load_system_prompt_s3_error(mock_s3_client):
     """Test system prompt loading with S3 error"""
     mock_s3_client.get_object.side_effect = Exception('S3 error')
-    
+
     with patch.dict('os.environ', {
         'SYSTEM_PROMPT_S3_BUCKET_NAME': 'test-bucket',
         'SYSTEM_PROMPT_S3_KEY': 'prompts/system.txt'
     }):
+        from index import load_system_prompt
         prompt = load_system_prompt()
-    
+
     assert prompt == "You are a helpful AI assistant."
 
 
-# Tests for SessionManager
+# Tests for load_tools_from_s3
 
-def test_session_manager_disabled():
-    """Test SessionManager with no bucket (disabled)"""
-    manager = SessionManager(bucket=None)
-    
-    assert not manager.enabled
-    assert manager.get_session('conn-123', 'user-123') == {'messages': [], 'metadata': {}}
+def test_load_tools_from_s3_empty_config():
+    """Test tool loading with empty config"""
+    from index import load_tools_from_s3
+    with patch.dict('os.environ', {'TOOLS_CONFIG': '[]'}):
+        tools = load_tools_from_s3()
+    assert tools == []
 
 
-def test_session_manager_get_session_success(mock_s3_client):
-    """Test successful session retrieval"""
-    session_data = {
-        'messages': [{'role': 'user', 'content': 'Hello'}],
-        'metadata': {'user_id': 'user-123'}
-    }
-    mock_s3_client.get_object.return_value = {
-        'Body': Mock(read=Mock(return_value=json.dumps(session_data).encode('utf-8')))
-    }
-    
-    manager = SessionManager(bucket='test-bucket')
-    result = manager.get_session('conn-123', 'user-123')
-    
-    assert result == session_data
-    mock_s3_client.get_object.assert_called_once_with(
-        Bucket='test-bucket',
-        Key='sessions/user-123/conn-123.json'
-    )
+def test_load_tools_from_s3_invalid_json():
+    """Test tool loading with invalid JSON config"""
+    from index import load_tools_from_s3
+    with patch.dict('os.environ', {'TOOLS_CONFIG': 'invalid'}):
+        tools = load_tools_from_s3()
+    assert tools == []
 
 
-def test_session_manager_get_session_not_found(mock_s3_client):
-    """Test session retrieval when session doesn't exist"""
-    mock_s3_client.get_object.side_effect = mock_s3_client.exceptions.NoSuchKey('Not found')
-    mock_s3_client.exceptions = type('Exceptions', (), {'NoSuchKey': Exception})()
-    
-    manager = SessionManager(bucket='test-bucket')
-    result = manager.get_session('conn-123', 'user-123')
-    
-    assert result == {'messages': [], 'metadata': {}}
+# Tests for SSE formatting
 
+def test_format_sse_data_only():
+    """Test SSE formatting with data only"""
+    from index import format_sse
+    result = format_sse('{"text": "hello"}')
+    assert result == 'data: {"text": "hello"}\n\n'
 
-def test_session_manager_save_session(mock_s3_client):
-    """Test session saving"""
-    session_data = {
-        'messages': [{'role': 'user', 'content': 'Hello'}],
-        'metadata': {'user_id': 'user-123'}
-    }
-    
-    manager = SessionManager(bucket='test-bucket')
-    manager.save_session('conn-123', 'user-123', session_data)
-    
-    mock_s3_client.put_object.assert_called_once()
-    call_args = mock_s3_client.put_object.call_args
-    assert call_args[1]['Bucket'] == 'test-bucket'
-    assert call_args[1]['Key'] == 'sessions/user-123/conn-123.json'
-    assert json.loads(call_args[1]['Body']) == session_data
-
-
-def test_session_manager_delete_session(mock_s3_client):
-    """Test session deletion"""
-    manager = SessionManager(bucket='test-bucket')
-    manager.delete_session('conn-123', 'user-123')
-    
-    mock_s3_client.delete_object.assert_called_once_with(
-        Bucket='test-bucket',
-        Key='sessions/user-123/conn-123.json'
-    )
 
+def test_format_sse_with_event():
+    """Test SSE formatting with event type"""
+    from index import format_sse
+    result = format_sse('{}', event='done')
+    assert result == 'event: done\ndata: {}\n\n'
 
-# Tests for ContextManager
-
-def test_context_manager_disabled():
-    """Test ContextManager with context disabled"""
-    with patch.dict('os.environ', {'CONTEXT_ENABLED': 'false'}):
-        manager = ContextManager()
-    
-    messages = [{'role': 'user', 'content': 'Hello'}]
-    assert manager.get_context(messages) == []
-    assert manager.add_message(messages, 'assistant', 'Hi') == messages
-
-
-def test_context_manager_sliding_window():
-    """Test SlidingWindow context strategy"""
-    manager = ContextManager(strategy='SlidingWindow', window_size=3)
-    
-    messages = [
-        {'role': 'user', 'content': 'Message 1'},
-        {'role': 'assistant', 'content': 'Response 1'},
-        {'role': 'user', 'content': 'Message 2'},
-        {'role': 'assistant', 'content': 'Response 2'},
-        {'role': 'user', 'content': 'Message 3'},
-    ]
-    
-    context = manager.get_context(messages)
-    assert len(context) == 3
-    assert context[0]['content'] == 'Message 2'
-
-
-def test_context_manager_null_strategy():
-    """Test Null context strategy"""
-    manager = ContextManager(strategy='Null')
-    
-    messages = [{'role': 'user', 'content': 'Hello'}]
-    assert manager.get_context(messages) == []
-
-
-def test_context_manager_add_message():
-    """Test adding message to context"""
-    manager = ContextManager()
-    
-    messages = []
-    messages = manager.add_message(messages, 'user', 'Hello')
-    messages = manager.add_message(messages, 'assistant', 'Hi there!')
-    
-    assert len(messages) == 2
-    assert messages[0] == {'role': 'user', 'content': 'Hello'}
-    assert messages[1] == {'role': 'assistant', 'content': 'Hi there!'}
-
-
-# Tests for send_to_connection
-
-def test_send_to_connection_success():
-    """Test successful message sending"""
-    mock_client = Mock()
-    
-    with patch('boto3.client', return_value=mock_client):
-        send_to_connection(
-            'https://test.execute-api.us-east-1.amazonaws.com/prod',
-            'conn-123',
-            {'type': 'chunk', 'content': 'Hello'}
-        )
-    
-    mock_client.post_to_connection.assert_called_once()
-    call_args = mock_client.post_to_connection.call_args
-    assert call_args[1]['ConnectionId'] == 'conn-123'
-    assert json.loads(call_args[1]['Data']) == {'type': 'chunk', 'content': 'Hello'}
-
-
-def test_send_to_connection_gone():
-    """Test sending to gone connection"""
-    mock_client = Mock()
-    mock_client.post_to_connection.side_effect = mock_client.exceptions.GoneException('Gone')
-    mock_client.exceptions = type('Exceptions', (), {'GoneException': Exception})()
-    
-    with patch('boto3.client', return_value=mock_client):
-        # Should not raise exception
-        send_to_connection(
-            'https://test.execute-api.us-east-1.amazonaws.com/prod',
-            'conn-123',
-            {'type': 'chunk', 'content': 'Hello'}
+
+# Tests for chat endpoint
+
+@pytest.mark.asyncio
+async def test_chat_creates_strands_session_manager():
+    """Test that chat creates Strands S3SessionManager when SESSION_BUCKET is set"""
+    from index import app
+    from httpx import AsyncClient, ASGITransport
+
+    mock_agent_instance = MagicMock()
+
+    # Mock stream_async to yield chunks and then return
+    async def mock_stream_async(msg):
+        yield {'data': 'Hello!'}
+
+    mock_agent_instance.stream_async = mock_stream_async
+
+    with patch('index.SESSION_BUCKET', 'test-bucket'), \
+         patch('index.StrandsS3SessionManager') as mock_session_cls, \
+         patch('index.SlidingWindowConversationManager') as mock_conv_cls, \
+         patch('index.Agent', return_value=mock_agent_instance) as mock_agent_cls:
+
+        transport = ASGITransport(app=app)
+        async with AsyncClient(transport=transport, base_url='http://test') as client:
+            response = await client.post('/chat', json={
+                'message': 'Hello',
+                'session_id': 'test-session-123',
+            })
+
+        assert response.status_code == 200
+
+        # Verify Strands S3SessionManager was created with correct params
+        mock_session_cls.assert_called_once_with(
+            session_id='test-session-123',
+            bucket_name='test-bucket',
         )
 
+        # Verify SlidingWindowConversationManager was created
+        mock_conv_cls.assert_called_once_with(window_size=20)
+
+        # Verify Agent was created with session_manager and conversation_manager
+        mock_agent_cls.assert_called_once()
+        call_kwargs = mock_agent_cls.call_args[1]
+        assert 'session_manager' in call_kwargs
+        assert 'conversation_manager' in call_kwargs
+        assert call_kwargs['callback_handler'] is None
+
 
-# Tests for handle_connect
+@pytest.mark.asyncio
+async def test_chat_no_session_bucket():
+    """Test that chat works without SESSION_BUCKET (no session persistence)"""
+    from index import app
+    from httpx import AsyncClient, ASGITransport
 
-def test_handle_connect_success(connect_event, mock_s3_client):
-    """Test successful connection handling"""
-    with patch.dict('os.environ', {'AUTH_TYPE': 'Cognito', 'SESSION_BUCKET': 'test-bucket'}):
-        response = handle_connect(connect_event)
-    
-    assert response['statusCode'] == 200
-    assert 'Connected' in response['body']
-    mock_s3_client.put_object.assert_called_once()
+    mock_agent_instance = MagicMock()
 
+    async def mock_stream_async(msg):
+        yield {'data': 'Hello!'}
 
-def test_handle_connect_anonymous(connect_event, mock_s3_client):
-    """Test connection handling without authentication"""
-    event = connect_event.copy()
-    del event['requestContext']['authorizer']
-    
-    with patch.dict('os.environ', {'AUTH_TYPE': 'None', 'SESSION_BUCKET': 'test-bucket'}):
-        response = handle_connect(event)
-    
-    assert response['statusCode'] == 200
+    mock_agent_instance.stream_async = mock_stream_async
 
+    with patch('index.SESSION_BUCKET', None), \
+         patch('index.Agent', return_value=mock_agent_instance) as mock_agent_cls:
 
-# Tests for handle_disconnect
+        transport = ASGITransport(app=app)
+        async with AsyncClient(transport=transport, base_url='http://test') as client:
+            response = await client.post('/chat', json={
+                'message': 'Hello',
+            })
 
-def test_handle_disconnect_success(disconnect_event, mock_s3_client):
-    """Test successful disconnection handling"""
-    with patch.dict('os.environ', {'AUTH_TYPE': 'Cognito', 'SESSION_BUCKET': 'test-bucket'}):
-        response = handle_disconnect(disconnect_event)
-    
-    assert response['statusCode'] == 200
-    assert 'Disconnected' in response['body']
-    mock_s3_client.delete_object.assert_called_once()
+        assert response.status_code == 200
 
+        # Agent should be called with session_manager=None
+        call_kwargs = mock_agent_cls.call_args[1]
+        assert call_kwargs['session_manager'] is None
 
-# Tests for handle_message
 
 @pytest.mark.asyncio
-async def test_handle_message_success(message_event, mock_s3_client):
-    """Test successful message handling with BidiAgent"""
-    # Mock session data
-    mock_s3_client.get_object.return_value = {
-        'Body': Mock(read=Mock(return_value=json.dumps({
-            'messages': [],
-            'metadata': {'user_id': 'user-123'}
-        }).encode('utf-8')))
-    }
-    
-    # Mock BidiAgent
-    mock_agent = AsyncMock()
-    mock_chunks = [
-        {'content': 'Hello '},
-        {'content': 'there! '},
-        {'content': 'How can I help?'}
-    ]
-    
-    async def mock_stream(*args, **kwargs):
-        for chunk in mock_chunks:
-            yield chunk
-    
-    mock_agent.stream = mock_stream
-    
-    # Mock send_to_connection
-    with patch('index.BidiAgent', return_value=mock_agent), \
-         patch('index.send_to_connection') as mock_send, \
-         patch.dict('os.environ', {
-             'SESSION_BUCKET': 'test-bucket',
-             'CONTEXT_ENABLED': 'true',
-             'AUTH_TYPE': 'Cognito'
-         }):
-        
-        response = await handle_message(message_event)
-    
-    assert response['statusCode'] == 200
-    assert 'Message processed' in response['body']
-    
-    # Verify chunks were sent
-    assert mock_send.call_count >= 3  # At least 3 chunks + complete message
-    
-    # Verify session was saved
-    mock_s3_client.put_object.assert_called_once()
+async def test_chat_generates_session_id():
+    """Test that chat generates a session_id when none provided"""
+    from index import app
+    from httpx import AsyncClient, ASGITransport
+
+    mock_agent_instance = MagicMock()
+
+    async def mock_stream_async(msg):
+        yield {'data': 'Hi'}
+
+    mock_agent_instance.stream_async = mock_stream_async
+
+    with patch('index.SESSION_BUCKET', None), \
+         patch('index.Agent', return_value=mock_agent_instance):
+
+        transport = ASGITransport(app=app)
+        async with AsyncClient(transport=transport, base_url='http://test') as client:
+            response = await client.post('/chat', json={
+                'message': 'Hello',
+            })
+
+        assert response.status_code == 200
+        # Parse SSE events to find metadata with session_id
+        lines = response.text.strip().split('\n')
+        metadata_data = None
+        for i, line in enumerate(lines):
+            if line.startswith('event: metadata'):
+                # Next non-empty line should be data:
+                for j in range(i + 1, len(lines)):
+                    if lines[j].startswith('data: '):
+                        metadata_data = json.loads(lines[j][6:])
+                        break
+                break
+        assert metadata_data is not None
+        assert 'session_id' in metadata_data
+        assert len(metadata_data['session_id']) > 0
 
 
 @pytest.mark.asyncio
-async def test_handle_message_invalid_json(message_event):
-    """Test message handling with invalid JSON"""
-    event = message_event.copy()
-    event['body'] = 'invalid json'
-    
-    with patch('index.send_to_connection') as mock_send:
-        response = await handle_message(event)
-    
-    assert response['statusCode'] == 400
-    assert 'Invalid JSON' in response['body']
-    
-    # Verify error was sent to client
-    mock_send.assert_called_once()
-    call_args = mock_send.call_args[0]
-    assert call_args[2]['type'] == 'error'
+async def test_chat_streams_sse_events():
+    """Test that chat streams SSE events correctly"""
+    from index import app
+    from httpx import AsyncClient, ASGITransport
+
+    mock_agent_instance = MagicMock()
+
+    async def mock_stream_async(msg):
+        yield {'data': 'Hello '}
+        yield {'data': 'world!'}
+
+    mock_agent_instance.stream_async = mock_stream_async
+
+    with patch('index.SESSION_BUCKET', None), \
+         patch('index.Agent', return_value=mock_agent_instance):
+
+        transport = ASGITransport(app=app)
+        async with AsyncClient(transport=transport, base_url='http://test') as client:
+            response = await client.post('/chat', json={
+                'message': 'Hello',
+                'session_id': 'test-session',
+            })
+
+        assert response.status_code == 200
+        body = response.text
+
+        # Should contain metadata, text chunks, and done events
+        assert 'event: metadata' in body
+        assert '"text": "Hello "' in body
+        assert '"text": "world!"' in body
+        assert 'event: done' in body
 
 
 @pytest.mark.asyncio
-async def test_handle_message_empty_message(message_event):
-    """Test message handling with empty message"""
-    event = message_event.copy()
-    event['body'] = json.dumps({'message': ''})
-    
-    with patch('index.send_to_connection') as mock_send:
-        response = await handle_message(event)
-    
-    assert response['statusCode'] == 400
-    assert 'Message is required' in response['body']
+async def test_chat_handles_agent_error():
+    """Test that chat handles agent errors gracefully"""
+    from index import app
+    from httpx import AsyncClient, ASGITransport
+
+    mock_agent_instance = MagicMock()
+
+    async def mock_stream_async(msg):
+        raise Exception('Agent error')
+        yield  # Make it an async generator
+
+    mock_agent_instance.stream_async = mock_stream_async
+
+    with patch('index.SESSION_BUCKET', None), \
+         patch('index.Agent', return_value=mock_agent_instance):
 
+        transport = ASGITransport(app=app)
+        async with AsyncClient(transport=transport, base_url='http://test') as client:
+            response = await client.post('/chat', json={
+                'message': 'Hello',
+                'session_id': 'test-session',
+            })
+
+        assert response.status_code == 200
+        body = response.text
+        assert 'event: error' in body
+        assert 'internal error' in body.lower() or 'error' in body.lower()
+
+
+# Tests for health endpoint
 
 @pytest.mark.asyncio
-async def test_handle_message_agent_error(message_event, mock_s3_client):
-    """Test message handling with agent error"""
-    # Mock session data
-    mock_s3_client.get_object.return_value = {
-        'Body': Mock(read=Mock(return_value=json.dumps({
-            'messages': [],
-            'metadata': {'user_id': 'user-123'}
-        }).encode('utf-8')))
-    }
-    
-    # Mock BidiAgent to raise error
-    mock_agent = AsyncMock()
-    mock_agent.stream.side_effect = Exception('Agent error')
-    
-    with patch('index.BidiAgent', return_value=mock_agent), \
-         patch('index.send_to_connection') as mock_send, \
-         patch.dict('os.environ', {
-             'SESSION_BUCKET': 'test-bucket',
-             'CONTEXT_ENABLED': 'true',
-             'AUTH_TYPE': 'Cognito'
-         }):
-        
-        response = await handle_message(message_event)
-    
-    assert response['statusCode'] == 500
-    assert 'Failed to process message' in response['body']
-    
-    # Verify error was sent to client
-    error_calls = [call for call in mock_send.call_args_list 
-                   if call[0][2].get('type') == 'error']
-    assert len(error_calls) > 0
-
-
-# Tests for main handler
-
-def test_handler_connect_route(connect_event, lambda_context, mock_s3_client):
-    """Test handler routing for $connect"""
-    with patch.dict('os.environ', {'SESSION_BUCKET': 'test-bucket'}):
-        response = handler(connect_event, lambda_context)
-    
-    assert response['statusCode'] == 200
-
-
-def test_handler_disconnect_route(disconnect_event, lambda_context, mock_s3_client):
-    """Test handler routing for $disconnect"""
-    with patch.dict('os.environ', {'SESSION_BUCKET': 'test-bucket'}):
-        response = handler(disconnect_event, lambda_context)
-    
-    assert response['statusCode'] == 200
-
-
-def test_handler_message_route(message_event, lambda_context, mock_s3_client):
-    """Test handler routing for message"""
-    # Mock session data
-    mock_s3_client.get_object.return_value = {
-        'Body': Mock(read=Mock(return_value=json.dumps({
-            'messages': [],
-            'metadata': {'user_id': 'user-123'}
-        }).encode('utf-8')))
-    }
-    
-    # Mock BidiAgent
-    mock_agent = AsyncMock()
-    
-    async def mock_stream(*args, **kwargs):
-        yield {'content': 'Hello'}
-    
-    mock_agent.stream = mock_stream
-    
-    with patch('index.BidiAgent', return_value=mock_agent), \
-         patch('index.send_to_connection'), \
-         patch.dict('os.environ', {
-             'SESSION_BUCKET': 'test-bucket',
-             'CONTEXT_ENABLED': 'true',
-             'AUTH_TYPE': 'Cognito'
-         }):
-        
-        response = handler(message_event, lambda_context)
-    
-    assert response['statusCode'] == 200
-
-
-def test_handler_unknown_route(lambda_context):
-    """Test handler with unknown route"""
-    event = {
-        'requestContext': {
-            'routeKey': 'unknown',
-            'connectionId': 'test-connection-123'
-        }
-    }
-    
-    response = handler(event, lambda_context)
-    
-    assert response['statusCode'] == 400
-    assert 'Unknown route' in response['body']
+async def test_health_endpoint():
+    """Test health check endpoint"""
+    from index import app
+    from httpx import AsyncClient, ASGITransport
+
+    transport = ASGITransport(app=app)
+    async with AsyncClient(transport=transport, base_url='http://test') as client:
+        response = await client.get('/health')
+
+    assert response.status_code == 200
+    assert response.json() == {'status': 'ok'}
+
+
+# Tests for handler function (Lambda compatibility)
+
+def test_handler_returns_200():
+    """Test Lambda handler fallback"""
+    from index import handler
+    result = handler({}, None)
+    assert result['statusCode'] == 200