mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
120 lines
4.0 KiB
Python
120 lines
4.0 KiB
Python
# Copyright 2025 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from unittest.mock import Mock
|
|
|
|
from google.adk.agents.base_agent import BaseAgent
|
|
from google.adk.agents.invocation_context import InvocationContext
|
|
from google.adk.events.event import Event
|
|
from google.adk.sessions.base_session_service import BaseSessionService
|
|
from google.adk.sessions.session import Session
|
|
import pytest
|
|
|
|
|
|
class TestInvocationContext:
|
|
"""Test suite for InvocationContext."""
|
|
|
|
@pytest.fixture
|
|
def mock_events(self):
|
|
"""Create mock events for testing."""
|
|
event1 = Mock(spec=Event)
|
|
event1.invocation_id = 'inv_1'
|
|
event1.branch = 'agent_1'
|
|
|
|
event2 = Mock(spec=Event)
|
|
event2.invocation_id = 'inv_1'
|
|
event2.branch = 'agent_2'
|
|
|
|
event3 = Mock(spec=Event)
|
|
event3.invocation_id = 'inv_2'
|
|
event3.branch = 'agent_1'
|
|
|
|
event4 = Mock(spec=Event)
|
|
event4.invocation_id = 'inv_2'
|
|
event4.branch = 'agent_2'
|
|
|
|
return [event1, event2, event3, event4]
|
|
|
|
@pytest.fixture
|
|
def mock_invocation_context(self, mock_events):
|
|
"""Create a mock invocation context for testing."""
|
|
ctx = InvocationContext(
|
|
session_service=Mock(spec=BaseSessionService),
|
|
agent=Mock(spec=BaseAgent),
|
|
invocation_id='inv_1',
|
|
branch='agent_1',
|
|
session=Mock(spec=Session, events=mock_events),
|
|
)
|
|
return ctx
|
|
|
|
def test_get_events_returns_all_events_by_default(
|
|
self, mock_invocation_context, mock_events
|
|
):
|
|
"""Tests that get_events returns all events when no filters are applied."""
|
|
events = mock_invocation_context.get_events()
|
|
assert events == mock_events
|
|
|
|
def test_get_events_filters_by_current_invocation(
|
|
self, mock_invocation_context, mock_events
|
|
):
|
|
"""Tests that get_events correctly filters by the current invocation."""
|
|
event1, event2, _, _ = mock_events
|
|
events = mock_invocation_context.get_events(current_invocation=True)
|
|
assert events == [event1, event2]
|
|
|
|
def test_get_events_filters_by_current_branch(
|
|
self, mock_invocation_context, mock_events
|
|
):
|
|
"""Tests that get_events correctly filters by the current branch."""
|
|
event1, _, event3, _ = mock_events
|
|
events = mock_invocation_context.get_events(current_branch=True)
|
|
assert events == [event1, event3]
|
|
|
|
def test_get_events_filters_by_invocation_and_branch(
|
|
self, mock_invocation_context, mock_events
|
|
):
|
|
"""Tests that get_events filters by invocation and branch."""
|
|
event1, _, _, _ = mock_events
|
|
events = mock_invocation_context.get_events(
|
|
current_invocation=True,
|
|
current_branch=True,
|
|
)
|
|
assert events == [event1]
|
|
|
|
def test_get_events_with_no_events_in_session(self, mock_invocation_context):
|
|
"""Tests get_events when the session has no events."""
|
|
mock_invocation_context.session.events = []
|
|
events = mock_invocation_context.get_events()
|
|
assert not events
|
|
|
|
def test_get_events_with_no_matching_events(self, mock_invocation_context):
|
|
"""Tests get_events when no events match the filters."""
|
|
mock_invocation_context.invocation_id = 'inv_3'
|
|
mock_invocation_context.branch = 'branch_C'
|
|
|
|
# Filter by invocation
|
|
events = mock_invocation_context.get_events(current_invocation=True)
|
|
assert not events
|
|
|
|
# Filter by branch
|
|
events = mock_invocation_context.get_events(current_branch=True)
|
|
assert not events
|
|
|
|
# Filter by both
|
|
events = mock_invocation_context.get_events(
|
|
current_invocation=True,
|
|
current_branch=True,
|
|
)
|
|
assert not events
|