278 lines
9.3 KiB
Python
278 lines
9.3 KiB
Python
"""Tests for db_migrate.py"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch, MagicMock
|
|
|
|
import pytest
|
|
|
|
# Add parent directory to path
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from db_migrate import Migration, MigrationManager
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_migrations_dir(tmp_path):
|
|
"""Create temporary migrations directory."""
|
|
migrations_dir = tmp_path / "migrations"
|
|
migrations_dir.mkdir()
|
|
return str(migrations_dir)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_mongo_client():
|
|
"""Mock MongoDB client."""
|
|
mock_client = MagicMock()
|
|
mock_db = MagicMock()
|
|
mock_client.get_default_database.return_value = mock_db
|
|
mock_client.server_info.return_value = {}
|
|
return mock_client, mock_db
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_postgres_conn():
|
|
"""Mock PostgreSQL connection."""
|
|
mock_conn = MagicMock()
|
|
mock_cursor = MagicMock()
|
|
mock_conn.cursor.return_value.__enter__.return_value = mock_cursor
|
|
return mock_conn, mock_cursor
|
|
|
|
|
|
class TestMigration:
|
|
"""Test Migration dataclass."""
|
|
|
|
def test_migration_creation(self):
|
|
"""Test creating migration object."""
|
|
migration = Migration(
|
|
id="20250101120000",
|
|
name="test_migration",
|
|
timestamp=datetime.now(),
|
|
database_type="mongodb"
|
|
)
|
|
|
|
assert migration.id == "20250101120000"
|
|
assert migration.name == "test_migration"
|
|
assert migration.database_type == "mongodb"
|
|
assert not migration.applied
|
|
|
|
|
|
class TestMigrationManager:
|
|
"""Test MigrationManager class."""
|
|
|
|
def test_init(self, temp_migrations_dir):
|
|
"""Test manager initialization."""
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
|
|
assert manager.db_type == "mongodb"
|
|
assert manager.connection_string == "mongodb://localhost"
|
|
assert Path(temp_migrations_dir).exists()
|
|
|
|
@patch('db_migrate.MongoClient')
|
|
def test_connect_mongodb(self, mock_client_class, temp_migrations_dir, mock_mongo_client):
|
|
"""Test MongoDB connection."""
|
|
mock_client, mock_db = mock_mongo_client
|
|
mock_client_class.return_value = mock_client
|
|
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
result = manager.connect()
|
|
|
|
assert result is True
|
|
assert manager.client == mock_client
|
|
assert manager.db == mock_db
|
|
|
|
@patch('db_migrate.psycopg2')
|
|
def test_connect_postgres(self, mock_psycopg2, temp_migrations_dir, mock_postgres_conn):
|
|
"""Test PostgreSQL connection."""
|
|
mock_conn, mock_cursor = mock_postgres_conn
|
|
mock_psycopg2.connect.return_value = mock_conn
|
|
|
|
manager = MigrationManager("postgres", "postgresql://localhost", temp_migrations_dir)
|
|
result = manager.connect()
|
|
|
|
assert result is True
|
|
assert manager.conn == mock_conn
|
|
|
|
def test_connect_unsupported_db(self, temp_migrations_dir):
|
|
"""Test connection with unsupported database type."""
|
|
manager = MigrationManager("unsupported", "connection_string", temp_migrations_dir)
|
|
result = manager.connect()
|
|
|
|
assert result is False
|
|
|
|
def test_generate_migration(self, temp_migrations_dir):
|
|
"""Test migration generation."""
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
migration = manager.generate_migration("test_migration")
|
|
|
|
assert migration is not None
|
|
assert migration.name == "test_migration"
|
|
|
|
# Check file was created
|
|
migration_files = list(Path(temp_migrations_dir).glob("*.json"))
|
|
assert len(migration_files) == 1
|
|
|
|
# Check file content
|
|
with open(migration_files[0]) as f:
|
|
data = json.load(f)
|
|
assert data["name"] == "test_migration"
|
|
assert data["database_type"] == "mongodb"
|
|
|
|
def test_generate_migration_dry_run(self, temp_migrations_dir):
|
|
"""Test migration generation in dry-run mode."""
|
|
manager = MigrationManager("postgres", "postgresql://localhost", temp_migrations_dir)
|
|
migration = manager.generate_migration("test_migration", dry_run=True)
|
|
|
|
assert migration is not None
|
|
|
|
# Check no file was created
|
|
migration_files = list(Path(temp_migrations_dir).glob("*.json"))
|
|
assert len(migration_files) == 0
|
|
|
|
def test_get_pending_migrations(self, temp_migrations_dir):
|
|
"""Test getting pending migrations."""
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
|
|
# Create test migration file
|
|
migration_data = {
|
|
"id": "20250101120000",
|
|
"name": "test_migration",
|
|
"timestamp": datetime.now().isoformat(),
|
|
"database_type": "mongodb",
|
|
"mongodb_operations": []
|
|
}
|
|
|
|
migration_file = Path(temp_migrations_dir) / "20250101120000_test.json"
|
|
with open(migration_file, "w") as f:
|
|
json.dump(migration_data, f)
|
|
|
|
# Mock database connection
|
|
with patch.object(manager, 'db', MagicMock()):
|
|
manager.db.migrations.find.return_value = []
|
|
|
|
pending = manager.get_pending_migrations()
|
|
|
|
assert len(pending) == 1
|
|
assert pending[0].id == "20250101120000"
|
|
assert pending[0].name == "test_migration"
|
|
|
|
@patch('db_migrate.MongoClient')
|
|
def test_apply_mongodb_migration(self, mock_client_class, temp_migrations_dir, mock_mongo_client):
|
|
"""Test applying MongoDB migration."""
|
|
mock_client, mock_db = mock_mongo_client
|
|
mock_client_class.return_value = mock_client
|
|
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
manager.connect()
|
|
|
|
migration = Migration(
|
|
id="20250101120000",
|
|
name="test_migration",
|
|
timestamp=datetime.now(),
|
|
database_type="mongodb",
|
|
mongodb_operations=[
|
|
{
|
|
"operation": "createIndex",
|
|
"collection": "users",
|
|
"index": {"email": 1},
|
|
"options": {}
|
|
}
|
|
]
|
|
)
|
|
|
|
result = manager.apply_migration(migration)
|
|
|
|
assert result is True
|
|
mock_db["users"].create_index.assert_called_once()
|
|
mock_db.migrations.insert_one.assert_called_once()
|
|
|
|
def test_apply_migration_dry_run(self, temp_migrations_dir):
|
|
"""Test applying migration in dry-run mode."""
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
|
|
migration = Migration(
|
|
id="20250101120000",
|
|
name="test_migration",
|
|
timestamp=datetime.now(),
|
|
database_type="mongodb",
|
|
mongodb_operations=[]
|
|
)
|
|
|
|
result = manager.apply_migration(migration, dry_run=True)
|
|
|
|
assert result is True
|
|
|
|
@patch('db_migrate.psycopg2')
|
|
def test_rollback_postgres_migration(self, mock_psycopg2, temp_migrations_dir, mock_postgres_conn):
|
|
"""Test rolling back PostgreSQL migration."""
|
|
mock_conn, mock_cursor = mock_postgres_conn
|
|
mock_psycopg2.connect.return_value = mock_conn
|
|
|
|
manager = MigrationManager("postgres", "postgresql://localhost", temp_migrations_dir)
|
|
manager.connect()
|
|
|
|
# Create migration file
|
|
migration_data = {
|
|
"id": "20250101120000",
|
|
"name": "test_migration",
|
|
"timestamp": datetime.now().isoformat(),
|
|
"database_type": "postgres",
|
|
"up_sql": "CREATE TABLE test (id INT);",
|
|
"down_sql": "DROP TABLE test;"
|
|
}
|
|
|
|
migration_file = Path(temp_migrations_dir) / "20250101120000_test.json"
|
|
with open(migration_file, "w") as f:
|
|
json.dump(migration_data, f)
|
|
|
|
result = manager.rollback_migration("20250101120000")
|
|
|
|
assert result is True
|
|
# Verify SQL was executed
|
|
assert mock_cursor.execute.call_count >= 1
|
|
|
|
def test_rollback_migration_not_found(self, temp_migrations_dir):
|
|
"""Test rollback with non-existent migration."""
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
|
|
result = manager.rollback_migration("99999999999999")
|
|
|
|
assert result is False
|
|
|
|
|
|
def test_migration_sorting(temp_migrations_dir):
|
|
"""Test that migrations are applied in correct order."""
|
|
manager = MigrationManager("mongodb", "mongodb://localhost", temp_migrations_dir)
|
|
|
|
# Create multiple migration files
|
|
for i in range(3):
|
|
migration_data = {
|
|
"id": f"2025010112000{i}",
|
|
"name": f"migration_{i}",
|
|
"timestamp": datetime.now().isoformat(),
|
|
"database_type": "mongodb",
|
|
"mongodb_operations": []
|
|
}
|
|
|
|
migration_file = Path(temp_migrations_dir) / f"2025010112000{i}_test.json"
|
|
with open(migration_file, "w") as f:
|
|
json.dump(migration_data, f)
|
|
|
|
with patch.object(manager, 'db', MagicMock()):
|
|
manager.db.migrations.find.return_value = []
|
|
|
|
pending = manager.get_pending_migrations()
|
|
|
|
# Check they're in order
|
|
assert len(pending) == 3
|
|
assert pending[0].id == "20250101120000"
|
|
assert pending[1].id == "20250101120001"
|
|
assert pending[2].id == "20250101120002"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|