|
- import os
- import asyncio
- import pytest
- from sqlmodel import Session, text, select
- from typing import Generator
-
- from autogenstudio.database import DatabaseManager
- from autogen_agentchat.agents import AssistantAgent
- from autogen_agentchat.teams import RoundRobinGroupChat
- from autogen_ext.models.openai import OpenAIChatCompletionClient
- from autogen_agentchat.conditions import TextMentionTermination
- from autogenstudio.datamodel.db import Team
-
-
- @pytest.fixture
- def test_db() -> Generator[DatabaseManager, None, None]:
- """Fixture for test database"""
- db_path = "test.db"
- db = DatabaseManager(f"sqlite:///{db_path}")
- db.reset_db()
- # Initialize database instead of create_db_and_tables
- db.initialize_database(auto_upgrade=False)
- yield db
- # Clean up
- asyncio.run(db.close())
- db.reset_db()
- try:
- if os.path.exists(db_path):
- os.remove(db_path)
- except Exception as e:
- print(f"Warning: Failed to remove test database file: {e}")
-
-
- @pytest.fixture
- def test_user() -> str:
- return "test_user@example.com"
-
-
- @pytest.fixture
- def sample_team(test_user: str) -> Team:
- """Create a sample team with proper config"""
- agent = AssistantAgent(
- name="weather_agent",
- model_client=OpenAIChatCompletionClient(
- model="gpt-4",
- ),
- )
-
- agent_team = RoundRobinGroupChat([agent], termination_condition=TextMentionTermination("TERMINATE"))
- team_component = agent_team.dump_component()
-
- return Team(
- user_id=test_user,
- component=team_component.model_dump(),
- )
-
-
- class TestDatabaseOperations:
- def test_basic_setup(self, test_db: DatabaseManager):
- """Test basic database setup and connection"""
- with Session(test_db.engine) as session:
- result = session.exec(text("SELECT 1")).first()
- assert result[0] == 1
- result = session.exec(select(1)).first()
- assert result == 1
-
- def test_basic_entity_creation(self, test_db: DatabaseManager, sample_team: Team):
- """Test creating all entity types with proper configs"""
- # Use upsert instead of raw session
- response = test_db.upsert(sample_team)
- assert response.status is True
-
- with Session(test_db.engine) as session:
- saved_team = session.get(Team, sample_team.id)
- assert saved_team is not None
-
- def test_upsert_operations(self, test_db: DatabaseManager, sample_team: Team):
- """Test upsert for both create and update scenarios"""
- # Test Create
- response = test_db.upsert(sample_team)
- assert response.status is True
- assert "Created Successfully" in response.message
-
- # Test Update
- team_id = sample_team.id
- sample_team.version = "0.0.2"
- response = test_db.upsert(sample_team)
- assert response.status is True
-
- # Verify Update
- result = test_db.get(Team, {"id": team_id})
- assert result.status is True
- assert result.data[0].version == "0.0.2"
-
- def test_delete_operations(self, test_db: DatabaseManager, sample_team: Team):
- """Test delete with various filters"""
- # First insert the model
- response = test_db.upsert(sample_team)
- assert response.status is True # Verify insert worked
-
- # Get the ID that was actually saved
- team_id = sample_team.id
-
- # Test deletion by id
- response = test_db.delete(Team, {"id": team_id})
- assert response.status is True
- assert "Deleted Successfully" in response.message
-
- # Verify deletion
- result = test_db.get(Team, {"id": team_id})
- assert len(result.data) == 0
-
- def test_initialize_database_scenarios(self):
- """Test different initialize_database parameters"""
- db_path = "test_init.db"
- db = DatabaseManager(f"sqlite:///{db_path}")
-
- try:
- # Test basic initialization
- response = db.initialize_database()
- assert response.status is True
-
- # Test with auto_upgrade
- response = db.initialize_database(auto_upgrade=True)
- assert response.status is True
-
- finally:
- asyncio.run(db.close())
- db.reset_db()
- if os.path.exists(db_path):
- os.remove(db_path)
|
