134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
"""
|
||
Test Database Management Script
|
||
This script helps create and manage the test database for unit tests.
|
||
"""
|
||
|
||
import psycopg2
|
||
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
|
||
import sys
|
||
|
||
# Test database configuration
|
||
TEST_DB_NAME = "test_ebook_db"
|
||
import os
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
TEST_DB_URL = os.getenv("TEST_DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/test_ebook_db")
|
||
|
||
def create_test_database():
|
||
"""Create test database if it doesn't exist"""
|
||
try:
|
||
# Connect to default postgres database to create test database
|
||
conn = psycopg2.connect(
|
||
host="localhost",
|
||
port="5432",
|
||
user="postgres",
|
||
password="postgres",
|
||
database="postgres"
|
||
)
|
||
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
|
||
cursor = conn.cursor()
|
||
|
||
# Check if test database exists
|
||
cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (TEST_DB_NAME,))
|
||
exists = cursor.fetchone()
|
||
|
||
if not exists:
|
||
cursor.execute(f"CREATE DATABASE {TEST_DB_NAME}")
|
||
print(f"✅ Created test database: {TEST_DB_NAME}")
|
||
else:
|
||
print(f"ℹ️ Test database {TEST_DB_NAME} already exists")
|
||
|
||
cursor.close()
|
||
conn.close()
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ Error creating test database: {e}")
|
||
return False
|
||
|
||
def drop_test_database():
|
||
"""Drop test database"""
|
||
try:
|
||
# Connect to default postgres database to drop test database
|
||
conn = psycopg2.connect(
|
||
host="localhost",
|
||
port="5432",
|
||
user="postgres",
|
||
password="postgres",
|
||
database="postgres"
|
||
)
|
||
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
|
||
cursor = conn.cursor()
|
||
|
||
# Terminate all connections to test database
|
||
cursor.execute(f"""
|
||
SELECT pg_terminate_backend(pid)
|
||
FROM pg_stat_activity
|
||
WHERE datname = '{TEST_DB_NAME}' AND pid <> pg_backend_pid()
|
||
""")
|
||
|
||
cursor.execute(f"DROP DATABASE IF EXISTS {TEST_DB_NAME}")
|
||
print(f"🗑️ Dropped test database: {TEST_DB_NAME}")
|
||
|
||
cursor.close()
|
||
conn.close()
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ Error dropping test database: {e}")
|
||
return False
|
||
|
||
def check_test_database():
|
||
"""Check if test database exists"""
|
||
try:
|
||
conn = psycopg2.connect(
|
||
host="localhost",
|
||
port="5432",
|
||
user="postgres",
|
||
password="postgres",
|
||
database="postgres"
|
||
)
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (TEST_DB_NAME,))
|
||
exists = cursor.fetchone()
|
||
|
||
cursor.close()
|
||
conn.close()
|
||
|
||
if exists:
|
||
print(f"✅ Test database {TEST_DB_NAME} exists")
|
||
return True
|
||
else:
|
||
print(f"❌ Test database {TEST_DB_NAME} does not exist")
|
||
return False
|
||
|
||
except Exception as e:
|
||
print(f"❌ Error checking test database: {e}")
|
||
return False
|
||
|
||
def main():
|
||
"""Main function to handle command line arguments"""
|
||
if len(sys.argv) < 2:
|
||
print("Usage: python manage_test_db.py [create|drop|check]")
|
||
print(" create - Create test database")
|
||
print(" drop - Drop test database")
|
||
print(" check - Check if test database exists")
|
||
return
|
||
|
||
command = sys.argv[1].lower()
|
||
|
||
if command == "create":
|
||
create_test_database()
|
||
elif command == "drop":
|
||
drop_test_database()
|
||
elif command == "check":
|
||
check_test_database()
|
||
else:
|
||
print(f"Unknown command: {command}")
|
||
print("Available commands: create, drop, check")
|
||
|
||
if __name__ == "__main__":
|
||
main() |