feat: Add get_events util function in invocation_context

PiperOrigin-RevId: 809111315

tests/unittests/agents/test_invocation_context.py

@@ -0,0 +1,119 @@
+# 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