Class MockPlugin

tests/unittests/test_runners.py:135–189

Mock plugin for unit testing.

class MockPlugin(BasePlugin):
  """Mock plugin for unit testing."""

  ON_USER_CALLBACK_MSG = (
      "Modified user message ON_USER_CALLBACK_MSG from MockPlugin"
  )
  ON_EVENT_CALLBACK_MSG = "Modified event ON_EVENT_CALLBACK_MSG from MockPlugin"
  ON_EVENT_CALLBACK_METADATA = {"plugin_key": "plugin_value"}

  def __init__(self):
    super().__init__(name="mock_plugin")
    self.enable_user_message_callback = False
    self.enable_event_callback = False
    self.user_content_seen_in_before_run_callback = None

  async def on_user_message_callback(
      self,
      *,
      invocation_context: InvocationContext,
      user_message: types.Content,
  ) -> Optional[types.Content]:
    if not self.enable_user_message_callback:
      return None
    return types.Content(
        role="model",
        parts=[types.Part(text=self.ON_USER_CALLBACK_MSG)],
    )

  async def before_run_callback(
      self,
      *,
      invocation_context: InvocationContext,
  ) -> None:
    self.user_content_seen_in_before_run_callback = (
        invocation_context.user_content
    )

  async def on_event_callback(
      self, *, invocation_context: InvocationContext, event: Event
  ) -> Optional[Event]:
    if not self.enable_event_callback:
      return None
    return Event(
        invocation_id="",
        author="",
        content=types.Content(
            parts=[
                types.Part(
                    text=self.ON_EVENT_CALLBACK_MSG,
                )
            ],
            role=event.content.role,
        ),
        custom_metadata=self.ON_EVENT_CALLBACK_METADATA,
    )


class TestRunnerFindAgentToRun: