Spaces:

sariskiat
/

donecase

Sleeping

donecase / tests /integration /test_agent_engine_app.py

mr.saris kiattithapanayong

update the code that demoed on saturday 22 nov

3d142aa 12 months ago

2.83 kB

# 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.

# mypy: disable-error-code="arg-type"

import logging

import pytest
from google.adk.events.event import Event

from rag_agent.agent_engine_app import AgentEngineApp


@pytest.fixture
def agent_app() -> AgentEngineApp:
    """Fixture to create and set up AgentEngineApp instance"""
    from rag_agent.agent_engine_app import agent_engine

    agent_engine.set_up()
    return agent_engine


@pytest.mark.asyncio
async def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
    """
    Integration test for the agent stream query functionality.
    Tests that the agent returns valid streaming responses.
    """
    # Create message and events for the async_stream_query
    message = "What's the weather in San Francisco?"
    events = []
    async for event in agent_app.async_stream_query(message=message, user_id="test"):
        events.append(event)
    assert len(events) > 0, "Expected at least one chunk in response"

    # Check for valid content in the response
    has_text_content = False
    for event in events:
        validated_event = Event.model_validate(event)
        content = validated_event.content
        if (
            content is not None
            and content.parts
            and any(part.text for part in content.parts)
        ):
            has_text_content = True
            break

    assert has_text_content, "Expected at least one event with text content"


def test_agent_feedback(agent_app: AgentEngineApp) -> None:
    """
    Integration test for the agent feedback functionality.
    Tests that feedback can be registered successfully.
    """
    feedback_data = {
        "score": 5,
        "text": "Great response!",
        "user_id": "test-user-456",
        "session_id": "test-session-456",
    }

    # Should not raise any exceptions
    agent_app.register_feedback(feedback_data)

    # Test invalid feedback
    with pytest.raises(ValueError):
        invalid_feedback = {
            "score": "invalid",  # Score must be numeric
            "text": "Bad feedback",
            "user_id": "test-user-789",
            "session_id": "test-session-789",
        }
        agent_app.register_feedback(invalid_feedback)

    logging.info("All assertions passed for agent feedback test")

	# 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.

	# mypy: disable-error-code="arg-type"

	import logging

	import pytest
	from google.adk.events.event import Event

	from rag_agent.agent_engine_app import AgentEngineApp


	@pytest.fixture
	def agent_app() -> AgentEngineApp:
	"""Fixture to create and set up AgentEngineApp instance"""
	from rag_agent.agent_engine_app import agent_engine

	agent_engine.set_up()
	return agent_engine


	@pytest.mark.asyncio
	async def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
	"""
	Integration test for the agent stream query functionality.
	Tests that the agent returns valid streaming responses.
	"""
	# Create message and events for the async_stream_query
	message = "What's the weather in San Francisco?"
	events = []
	async for event in agent_app.async_stream_query(message=message, user_id="test"):
	events.append(event)
	assert len(events) > 0, "Expected at least one chunk in response"

	# Check for valid content in the response
	has_text_content = False
	for event in events:
	validated_event = Event.model_validate(event)
	content = validated_event.content
	if (
	content is not None
	and content.parts
	and any(part.text for part in content.parts)
	):
	has_text_content = True
	break

	assert has_text_content, "Expected at least one event with text content"


	def test_agent_feedback(agent_app: AgentEngineApp) -> None:
	"""
	Integration test for the agent feedback functionality.
	Tests that feedback can be registered successfully.
	"""
	feedback_data = {
	"score": 5,
	"text": "Great response!",
	"user_id": "test-user-456",
	"session_id": "test-session-456",
	}

	# Should not raise any exceptions
	agent_app.register_feedback(feedback_data)

	# Test invalid feedback
	with pytest.raises(ValueError):
	invalid_feedback = {
	"score": "invalid", # Score must be numeric
	"text": "Bad feedback",
	"user_id": "test-user-789",
	"session_id": "test-session-789",
	}
	agent_app.register_feedback(invalid_feedback)

	logging.info("All assertions passed for agent feedback test")