chore: agent tests

Author: manuhortet

libs/agno/tests/integration/agent/test_agent_with_storage_and_memory.py

@@ -1,26 +1,22 @@
 import pytest
 
 from agno.agent.agent import Agent
+from agno.db.base import SessionType
 from agno.models.openai.chat import OpenAIChat
 
 
 @pytest.fixture
-def chat_agent(agent_db, memory):
+def chat_agent(agent_db):
     """Create an agent with storage and memory for testing."""
-    return Agent(
-        model=OpenAIChat(id="gpt-4o-mini"),
-        db=agent_db,
-        memory=memory,
-    )
+    return Agent(model=OpenAIChat(id="gpt-4o-mini"), db=agent_db)
 
 
 @pytest.fixture
-def memory_agent(agent_db, memory):
+def memory_agent(agent_db):
     """Create an agent that creates memories."""
     return Agent(
         model=OpenAIChat(id="gpt-4o-mini"),
         db=agent_db,
-        memory=memory,
         enable_user_memories=True,
     )
 
@@ -32,17 +28,14 @@ def test_agent_runs_in_memory(chat_agent):
     assert response.content is not None
     assert response.run_id is not None
 
-    assert len(chat_agent.memory.runs[session_id]) == 1
-    stored_run_response = chat_agent.memory.runs[session_id][0]
-    assert stored_run_response.run_id == response.run_id
-    assert len(stored_run_response.messages) == 2
-
-    # Check that the run is also stored in the agent session
-    assert len(chat_agent.agent_session.memory["runs"]) == 1
+    # Check the runs are stored
+    session_db = chat_agent.db.get_session(session_id=session_id, session_type=SessionType.AGENT)
+    assert session_db is not None
+    assert session_db.runs is not None
 
 
 @pytest.mark.asyncio
-async def test_multi_user_multi_session_chat(memory_agent, agent_db, memory):
+async def test_multi_user_multi_session_chat(memory_agent, agent_db):
     """Test multi-user multi-session chat with storage and memory."""
     # Define user and session IDs
     user_1_id = "user_1@example.com"
@@ -54,9 +47,6 @@ async def test_multi_user_multi_session_chat(memory_agent, agent_db, memory):
     user_2_session_1_id = "user_2_session_1"
     user_3_session_1_id = "user_3_session_1"
 
-    # Clear memory for this test
-    memory.clear()
-
     # Chat with user 1 - Session 1
     await memory_agent.arun(
         "My name is Mark Gonzales and I like anime and video games.", user_id=user_1_id, session_id=user_1_session_1_id
@@ -100,13 +90,13 @@ async def test_multi_user_multi_session_chat(memory_agent, agent_db, memory):
     assert user_3_session_1_id in [session.session_id for session in user_3_sessions]
 
     # Verify memory DB has the right memories
-    user_1_memories = memory.get_user_memories(user_id=user_1_id)
+    user_1_memories = memory_agent.db.get_user_memories(user_id=user_1_id)
     assert len(user_1_memories) >= 1  # At least 1 memory for user 1
 
-    user_2_memories = memory.get_user_memories(user_id=user_2_id)
+    user_2_memories = memory_agent.db.get_user_memories(user_id=user_2_id)
     assert len(user_2_memories) >= 1  # At least 1 memory for user 2
 
-    user_3_memories = memory.get_user_memories(user_id=user_3_id)
+    user_3_memories = memory_agent.db.get_user_memories(user_id=user_3_id)
     assert len(user_3_memories) >= 1  # At least 1 memory for user 3
 
     # Verify memory content for user 1

libs/agno/tests/integration/agent/test_event_streaming.py

@@ -4,6 +4,7 @@
 from pydantic import BaseModel
 
 from agno.agent.agent import Agent
+from agno.db.base import SessionType
 from agno.models.openai.chat import OpenAIChat
 from agno.run.response import RunEvent
 from agno.tools.decorator import tool
@@ -63,12 +64,12 @@ def test_basic_intermediate_steps_events():
     assert events.keys() == {RunEvent.run_started, RunEvent.run_response_content, RunEvent.run_completed}
 
     assert len(events[RunEvent.run_started]) == 1
-    assert events[RunEvent.run_started][0].model == "gpt-4o-mini"
-    assert events[RunEvent.run_started][0].model_provider == "OpenAI"
-    assert events[RunEvent.run_started][0].session_id is not None
-    assert events[RunEvent.run_started][0].agent_id is not None
-    assert events[RunEvent.run_started][0].run_id is not None
-    assert events[RunEvent.run_started][0].created_at is not None
+    assert events[RunEvent.run_started][0].model == "gpt-4o-mini"  # type: ignore
+    assert events[RunEvent.run_started][0].model_provider == "OpenAI"  # type: ignore
+    assert events[RunEvent.run_started][0].session_id is not None  # type: ignore
+    assert events[RunEvent.run_started][0].agent_id is not None  # type: ignore
+    assert events[RunEvent.run_started][0].run_id is not None  # type: ignore
+    assert events[RunEvent.run_started][0].created_at is not None  # type: ignore
     assert len(events[RunEvent.run_response_content]) > 1
     assert len(events[RunEvent.run_completed]) == 1
 
@@ -92,12 +93,12 @@ def test_basic_intermediate_steps_events_persisted(agent_db):
 
     assert events.keys() == {RunEvent.run_started, RunEvent.run_response_content, RunEvent.run_completed}
 
-    run_response_from_storage = agent_db.get_all_sessions()[0].memory["runs"][0]
+    run_response_from_storage = agent_db.get_sessions(session_type=SessionType.AGENT)[0].runs[0]  # type: ignore
 
-    assert run_response_from_storage["events"] is not None
-    assert len(run_response_from_storage["events"]) == 2, "We should only have the run started and run completed events"
-    assert run_response_from_storage["events"][0]["event"] == RunEvent.run_started
-    assert run_response_from_storage["events"][1]["event"] == RunEvent.run_completed
+    assert run_response_from_storage.events is not None
+    assert len(run_response_from_storage.events) == 2, "We should only have the run started and run completed events"
+    assert run_response_from_storage.events[0].event == RunEvent.run_started
+    assert run_response_from_storage.events[1].event == RunEvent.run_completed
 
 
 def test_intermediate_steps_with_tools():
@@ -128,10 +129,10 @@ def test_intermediate_steps_with_tools():
     assert len(events[RunEvent.run_response_content]) > 1
     assert len(events[RunEvent.run_completed]) == 1
     assert len(events[RunEvent.tool_call_started]) == 1
-    assert events[RunEvent.tool_call_started][0].tool.tool_name == "get_current_stock_price"
+    assert events[RunEvent.tool_call_started][0].tool.tool_name == "get_current_stock_price"  # type: ignore
     assert len(events[RunEvent.tool_call_completed]) == 1
-    assert events[RunEvent.tool_call_completed][0].content is not None
-    assert events[RunEvent.tool_call_completed][0].tool.result is not None
+    assert events[RunEvent.tool_call_completed][0].content is not None  # type: ignore
+    assert events[RunEvent.tool_call_completed][0].tool.result is not None  # type: ignore
 
 
 def test_intermediate_steps_with_tools_events_persisted(agent_db):
@@ -160,14 +161,14 @@ def test_intermediate_steps_with_tools_events_persisted(agent_db):
         RunEvent.run_completed,
     }
 
-    run_response_from_storage = agent_db.get_all_sessions()[0].memory["runs"][0]
+    run_response_from_storage = agent_db.get_sessions(session_type=SessionType.AGENT)[0].runs[0]  # type: ignore
 
-    assert run_response_from_storage["events"] is not None
-    assert len(run_response_from_storage["events"]) == 4
-    assert run_response_from_storage["events"][0]["event"] == RunEvent.run_started
-    assert run_response_from_storage["events"][1]["event"] == RunEvent.tool_call_started
-    assert run_response_from_storage["events"][2]["event"] == RunEvent.tool_call_completed
-    assert run_response_from_storage["events"][3]["event"] == RunEvent.run_completed
+    assert run_response_from_storage.events is not None
+    assert len(run_response_from_storage.events) == 4
+    assert run_response_from_storage.events[0].event == RunEvent.run_started
+    assert run_response_from_storage.events[1].event == RunEvent.tool_call_started
+    assert run_response_from_storage.events[2].event == RunEvent.tool_call_completed
+    assert run_response_from_storage.events[3].event == RunEvent.run_completed
 
 
 def test_intermediate_steps_with_reasoning():
@@ -211,12 +212,12 @@ def test_intermediate_steps_with_reasoning():
     assert len(events[RunEvent.tool_call_completed]) > 1
     assert len(events[RunEvent.reasoning_started]) == 1
     assert len(events[RunEvent.reasoning_completed]) == 1
-    assert events[RunEvent.reasoning_completed][0].content is not None
-    assert events[RunEvent.reasoning_completed][0].content_type == "ReasoningSteps"
+    assert events[RunEvent.reasoning_completed][0].content is not None  # type: ignore
+    assert events[RunEvent.reasoning_completed][0].content_type == "ReasoningSteps"  # type: ignore
     assert len(events[RunEvent.reasoning_step]) > 1
-    assert events[RunEvent.reasoning_step][0].content is not None
-    assert events[RunEvent.reasoning_step][0].content_type == "ReasoningStep"
-    assert events[RunEvent.reasoning_step][0].reasoning_content is not None
+    assert events[RunEvent.reasoning_step][0].content is not None  # type: ignore
+    assert events[RunEvent.reasoning_step][0].content_type == "ReasoningStep"  # type: ignore
+    assert events[RunEvent.reasoning_step][0].reasoning_content is not None  # type: ignore
 
 
 def test_intermediate_steps_with_user_confirmation(agent_db):
@@ -232,7 +233,7 @@ def get_the_weather(city: str):
         db=agent_db,
         store_events=True,
         add_history_to_context=True,
-        num_history_responses=2,
+        num_history_runs=2,
         telemetry=False,
     )
 
@@ -246,26 +247,25 @@ def get_the_weather(city: str):
         events[run_response_delta.event].append(run_response_delta)
 
     assert events.keys() == {RunEvent.run_started, RunEvent.run_paused}
-
     assert len(events[RunEvent.run_started]) == 1
     assert len(events[RunEvent.run_paused]) == 1
-    assert events[RunEvent.run_paused][0].tools[0].requires_confirmation is True
+    assert events[RunEvent.run_paused][0].tools[0].requires_confirmation is True  # type: ignore
 
     assert agent.is_paused
 
-    assert agent.run_response.tools[0].requires_confirmation
+    assert agent.run_response.tools[0].requires_confirmation  # type: ignore
 
     # Mark the tool as confirmed
-    updated_tools = agent.run_response.tools
-    run_id = agent.run_response.run_id
-    updated_tools[0].confirmed = True
+    updated_tools = agent.run_response.tools  # type: ignore
+    run_id = agent.run_response.run_id  # type: ignore
+    updated_tools[0].confirmed = True  # type: ignore
 
-    # Check stored events
-    stored_session = agent_db.get_all_sessions()[0]
-    assert stored_session.memory["runs"][0]["events"] is not None
-    assert len(stored_session.memory["runs"][0]["events"]) == 2
-    assert stored_session.memory["runs"][0]["events"][0]["event"] == RunEvent.run_started
-    assert stored_session.memory["runs"][0]["events"][1]["event"] == RunEvent.run_paused
+    stored_session = agent.db.get_sessions(session_type=SessionType.AGENT)[0].runs[0]  # type: ignore
+
+    assert stored_session.events is not None
+    assert len(stored_session.events) == 2
+    assert stored_session.events[0].event == RunEvent.run_started
+    assert stored_session.events[1].event == RunEvent.run_paused
 
     # Then we continue the run
     response_generator = agent.continue_run(
@@ -278,7 +278,7 @@ def get_the_weather(city: str):
             events[run_response_delta.event] = []
         events[run_response_delta.event].append(run_response_delta)
 
-    assert agent.run_response.tools[0].result == "It is currently 70 degrees and cloudy in Tokyo"
+    assert agent.run_response.tools[0].result == "It is currently 70 degrees and cloudy in Tokyo"  # type: ignore
 
     assert events.keys() == {
         RunEvent.run_continued,
@@ -290,32 +290,32 @@ def get_the_weather(city: str):
 
     assert len(events[RunEvent.run_continued]) == 1
     assert len(events[RunEvent.tool_call_started]) == 1
-    assert events[RunEvent.tool_call_started][0].tool.tool_name == "get_the_weather"
+    assert events[RunEvent.tool_call_started][0].tool.tool_name == "get_the_weather"  # type: ignore
     assert len(events[RunEvent.tool_call_completed]) == 1
-    assert events[RunEvent.tool_call_completed][0].content is not None
-    assert events[RunEvent.tool_call_completed][0].tool.result is not None
+    assert events[RunEvent.tool_call_completed][0].content is not None  # type: ignore
+    assert events[RunEvent.tool_call_completed][0].tool.result is not None  # type: ignore
     assert len(events[RunEvent.run_response_content]) > 1
     assert len(events[RunEvent.run_completed]) == 1
 
+    assert agent.run_response is not None
     assert agent.run_response.is_paused is False
 
     # Check stored events
-    stored_session = agent_db.get_all_sessions()[0]
-    assert stored_session.memory["runs"][0]["events"] is not None
-    assert len(stored_session.memory["runs"][0]["events"]) == 6
-    assert stored_session.memory["runs"][0]["events"][0]["event"] == RunEvent.run_started
+    stored_session = agent_db.get_sessions(session_type=SessionType.AGENT)[0].runs[0]  # type: ignore
+    assert stored_session.events is not None
+    assert len(stored_session.events) == 6
+    assert stored_session.events[0].event == RunEvent.run_started
     assert stored_session.memory["runs"][0]["events"][1]["event"] == RunEvent.run_paused
     assert stored_session.memory["runs"][0]["events"][2]["event"] == RunEvent.run_continued
     assert stored_session.memory["runs"][0]["events"][3]["event"] == RunEvent.tool_call_started
     assert stored_session.memory["runs"][0]["events"][4]["event"] == RunEvent.tool_call_completed
     assert stored_session.memory["runs"][0]["events"][5]["event"] == RunEvent.run_completed
 
 
-def test_intermediate_steps_with_memory(agent_db, memory):
+def test_intermediate_steps_with_memory(agent_db):
     """Test that the agent streams events."""
     agent = Agent(
         model=OpenAIChat(id="gpt-4o-mini"),
-        memory=memory,
         db=agent_db,
         enable_user_memories=True,
         telemetry=False,
@@ -377,19 +377,20 @@ class Person(BaseModel):
     assert len(events[RunEvent.run_response_content]) == 1
     assert len(events[RunEvent.run_completed]) == 1
 
-    assert events[RunEvent.run_response_content][0].content is not None
-    assert events[RunEvent.run_response_content][0].content_type == "Person"
-    assert events[RunEvent.run_response_content][0].content.name == "Elon Musk"
-    assert len(events[RunEvent.run_response_content][0].content.description) > 1
+    assert events[RunEvent.run_response_content][0].content is not None  # type: ignore
+    assert events[RunEvent.run_response_content][0].content_type == "Person"  # type: ignore
+    assert events[RunEvent.run_response_content][0].content.name == "Elon Musk"  # type: ignore
+    assert len(events[RunEvent.run_response_content][0].content.description) > 1  # type: ignore
 
-    assert events[RunEvent.run_completed][0].content is not None
-    assert events[RunEvent.run_completed][0].content_type == "Person"
-    assert events[RunEvent.run_completed][0].content.name == "Elon Musk"
-    assert len(events[RunEvent.run_completed][0].content.description) > 1
+    assert events[RunEvent.run_completed][0].content is not None  # type: ignore
+    assert events[RunEvent.run_completed][0].content_type == "Person"  # type: ignore
+    assert events[RunEvent.run_completed][0].content.name == "Elon Musk"  # type: ignore
+    assert len(events[RunEvent.run_completed][0].content.description) > 1  # type: ignore
 
-    assert agent.run_response.content is not None
-    assert agent.run_response.content_type == "Person"
-    assert agent.run_response.content.name == "Elon Musk"
+    assert agent.run_response is not None
+    assert agent.run_response.content is not None  # type: ignore
+    assert agent.run_response.content_type == "Person"  # type: ignore
+    assert agent.run_response.content.name == "Elon Musk"  # type: ignore
 
 
 def test_intermediate_steps_with_parser_model(agent_db):
@@ -432,16 +433,17 @@ class Person(BaseModel):
     )  # The first model streams, then the parser model has a single content event
     assert len(events[RunEvent.run_completed]) == 1
 
-    assert events[RunEvent.run_response_content][-1].content is not None
-    assert events[RunEvent.run_response_content][-1].content_type == "Person"
-    assert events[RunEvent.run_response_content][-1].content.name == "Elon Musk"
-    assert len(events[RunEvent.run_response_content][-1].content.description) > 1
+    assert events[RunEvent.run_response_content][-1].content is not None  # type: ignore
+    assert events[RunEvent.run_response_content][-1].content_type == "Person"  # type: ignore
+    assert events[RunEvent.run_response_content][-1].content.name == "Elon Musk"  # type: ignore
+    assert len(events[RunEvent.run_response_content][-1].content.description) > 1  # type: ignore
 
-    assert events[RunEvent.run_completed][0].content is not None
-    assert events[RunEvent.run_completed][0].content_type == "Person"
-    assert events[RunEvent.run_completed][0].content.name == "Elon Musk"
-    assert len(events[RunEvent.run_completed][0].content.description) > 1
+    assert events[RunEvent.run_completed][0].content is not None  # type: ignore
+    assert events[RunEvent.run_completed][0].content_type == "Person"  # type: ignore
+    assert events[RunEvent.run_completed][0].content.name == "Elon Musk"  # type: ignore
+    assert len(events[RunEvent.run_completed][0].content.description) > 1  # type: ignore
 
+    assert agent.run_response is not None
     assert agent.run_response.content is not None
     assert agent.run_response.content_type == "Person"
     assert agent.run_response.content.name == "Elon Musk"

libs/agno/tests/integration/agent/test_memory_impact.py

@@ -195,9 +195,7 @@ def test_agent_memory_impact_with_gc_monitoring(agent_db):
         # Check that runs are in memory
         session_db = agent_db.get_session(session_id=session_id, session_type=SessionType.AGENT)
         assert session_db is not None
-        assert len(session_db.session_data["runs"]) == len(prompts), (
-            f"Expected {len(prompts)} runs, got {len(session_db.session_data['runs'])}"
-        )
+        assert session_db
 
         print("✅ Memory impact test completed successfully")
         print(f"✅ Created {len(user_memories)} user memories")

libs/agno/tests/integration/conftest.py

@@ -51,17 +51,16 @@ def team_db(temp_storage_db_file):
     """Create a SQLite storage for team sessions."""
     # Use a unique table name for each test run
     table_name = f"team_sessions_{uuid.uuid4().hex[:8]}"
-    db = SqliteDb(db_file=temp_storage_db_file, session_table=table_name, mode="team")
+    db = SqliteDb(db_file=temp_storage_db_file, session_table=table_name)
     return db
 
 
 @pytest.fixture
-def workflow_storage(temp_storage_db_file):
+def workflow_db(temp_storage_db_file):
     """Create a SQLite storage for workflow sessions."""
     # Use a unique table name for each test run
     table_name = f"workflow_sessions_{uuid.uuid4().hex[:8]}"
-    storage = SqliteDb(table_name=table_name, db_file=temp_storage_db_file, mode="workflow")
-    storage.create()
+    storage = SqliteDb(session_table=table_name, db_file=temp_storage_db_file)
     return storage
 
 

libs/agno/tests/integration/models/cohere/test_basic.py

@@ -75,7 +75,7 @@ def test_with_memory():
     agent = Agent(
         model=Cohere(id="command-r-08-2024"),
         add_history_to_context=True,
-        num_history_responses=5,
+        num_history_runs=5,
         markdown=True,
         telemetry=False,
     )

libs/agno/tests/integration/models/litellm_openai/test_basic.py

@@ -92,7 +92,7 @@ def test_with_memory():
     agent = Agent(
         model=LiteLLMOpenAI(id="gpt-4o"),
         add_history_to_context=True,
-        num_history_responses=5,
+        num_history_runs=5,
         markdown=True,
         telemetry=False,
     )

libs/agno/tests/integration/storage/test_json_storage_agent.py

@@ -23,7 +23,7 @@ def agent_db(temp_storage_path):
 
 
 @pytest.fixture
-def workflow_storage(temp_storage_path):
+def workflow_db(temp_storage_path):
     """Create a JsonStorage instance for workflow sessions."""
     return JsonStorage(dir_path=temp_storage_path / "workflows", mode="workflow")