mirror of
https://github.com/Xe138/AI-Trader.git
synced 2026-04-01 17:17:24 -04:00
feat: add dev database initialization and cleanup functions
This commit is contained in:
@@ -213,6 +213,57 @@ def initialize_database(db_path: str = "data/jobs.db") -> None:
|
|||||||
conn.close()
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
def initialize_dev_database(db_path: str = "data/trading_dev.db") -> None:
|
||||||
|
"""
|
||||||
|
Initialize dev database with clean schema
|
||||||
|
|
||||||
|
Deletes and recreates dev database unless PRESERVE_DEV_DATA=true.
|
||||||
|
Used at startup in DEV mode to ensure clean testing environment.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
db_path: Path to dev database file
|
||||||
|
"""
|
||||||
|
from tools.deployment_config import should_preserve_dev_data
|
||||||
|
|
||||||
|
if should_preserve_dev_data():
|
||||||
|
print(f"ℹ️ PRESERVE_DEV_DATA=true, keeping existing dev database: {db_path}")
|
||||||
|
# Ensure schema exists even if preserving data
|
||||||
|
if not Path(db_path).exists():
|
||||||
|
print(f"📁 Dev database doesn't exist, creating: {db_path}")
|
||||||
|
initialize_database(db_path)
|
||||||
|
return
|
||||||
|
|
||||||
|
# Delete existing dev database
|
||||||
|
if Path(db_path).exists():
|
||||||
|
print(f"🗑️ Removing existing dev database: {db_path}")
|
||||||
|
Path(db_path).unlink()
|
||||||
|
|
||||||
|
# Create fresh dev database
|
||||||
|
print(f"📁 Creating fresh dev database: {db_path}")
|
||||||
|
initialize_database(db_path)
|
||||||
|
|
||||||
|
|
||||||
|
def cleanup_dev_database(db_path: str = "data/trading_dev.db", data_path: str = "./data/dev_agent_data") -> None:
|
||||||
|
"""
|
||||||
|
Cleanup dev database and data files
|
||||||
|
|
||||||
|
Args:
|
||||||
|
db_path: Path to dev database file
|
||||||
|
data_path: Path to dev data directory
|
||||||
|
"""
|
||||||
|
import shutil
|
||||||
|
|
||||||
|
# Remove dev database
|
||||||
|
if Path(db_path).exists():
|
||||||
|
print(f"🗑️ Removing dev database: {db_path}")
|
||||||
|
Path(db_path).unlink()
|
||||||
|
|
||||||
|
# Remove dev data directory
|
||||||
|
if Path(data_path).exists():
|
||||||
|
print(f"🗑️ Removing dev data directory: {data_path}")
|
||||||
|
shutil.rmtree(data_path)
|
||||||
|
|
||||||
|
|
||||||
def _migrate_schema(cursor: sqlite3.Cursor) -> None:
|
def _migrate_schema(cursor: sqlite3.Cursor) -> None:
|
||||||
"""Migrate existing database schema to latest version."""
|
"""Migrate existing database schema to latest version."""
|
||||||
# Check if positions table exists and has simulation_run_id column
|
# Check if positions table exists and has simulation_run_id column
|
||||||
|
|||||||
84
tests/unit/test_dev_database.py
Normal file
84
tests/unit/test_dev_database.py
Normal file
@@ -0,0 +1,84 @@
|
|||||||
|
import os
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from api.database import initialize_dev_database, cleanup_dev_database
|
||||||
|
|
||||||
|
|
||||||
|
def test_initialize_dev_database_creates_fresh_db(tmp_path):
|
||||||
|
"""Test dev database initialization creates clean schema"""
|
||||||
|
db_path = str(tmp_path / "test_dev.db")
|
||||||
|
|
||||||
|
# Create initial database with some data
|
||||||
|
from api.database import get_db_connection, initialize_database
|
||||||
|
initialize_database(db_path)
|
||||||
|
conn = get_db_connection(db_path)
|
||||||
|
conn.execute("INSERT INTO jobs (job_id, config_path, status, date_range, models, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||||
|
("test-job", "config.json", "completed", "2025-01-01:2025-01-31", '["model1"]', "2025-01-01T00:00:00"))
|
||||||
|
conn.commit()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
# Verify data exists
|
||||||
|
conn = get_db_connection(db_path)
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute("SELECT COUNT(*) FROM jobs")
|
||||||
|
assert cursor.fetchone()[0] == 1
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
# Initialize dev database (should reset)
|
||||||
|
initialize_dev_database(db_path)
|
||||||
|
|
||||||
|
# Verify data is cleared
|
||||||
|
conn = get_db_connection(db_path)
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute("SELECT COUNT(*) FROM jobs")
|
||||||
|
assert cursor.fetchone()[0] == 0
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_dev_database_removes_files(tmp_path):
|
||||||
|
"""Test dev cleanup removes database and data files"""
|
||||||
|
# Setup dev files
|
||||||
|
db_path = str(tmp_path / "test_dev.db")
|
||||||
|
data_path = str(tmp_path / "dev_agent_data")
|
||||||
|
|
||||||
|
Path(db_path).touch()
|
||||||
|
Path(data_path).mkdir(parents=True, exist_ok=True)
|
||||||
|
(Path(data_path) / "test_file.jsonl").touch()
|
||||||
|
|
||||||
|
# Verify files exist
|
||||||
|
assert Path(db_path).exists()
|
||||||
|
assert Path(data_path).exists()
|
||||||
|
|
||||||
|
# Cleanup
|
||||||
|
cleanup_dev_database(db_path, data_path)
|
||||||
|
|
||||||
|
# Verify files removed
|
||||||
|
assert not Path(db_path).exists()
|
||||||
|
assert not Path(data_path).exists()
|
||||||
|
|
||||||
|
|
||||||
|
def test_initialize_dev_respects_preserve_flag(tmp_path):
|
||||||
|
"""Test that PRESERVE_DEV_DATA flag prevents cleanup"""
|
||||||
|
os.environ["PRESERVE_DEV_DATA"] = "true"
|
||||||
|
db_path = str(tmp_path / "test_dev.db")
|
||||||
|
|
||||||
|
# Create database with data
|
||||||
|
from api.database import get_db_connection, initialize_database
|
||||||
|
initialize_database(db_path)
|
||||||
|
conn = get_db_connection(db_path)
|
||||||
|
conn.execute("INSERT INTO jobs (job_id, config_path, status, date_range, models, created_at) VALUES (?, ?, ?, ?, ?, ?)",
|
||||||
|
("test-job", "config.json", "completed", "2025-01-01:2025-01-31", '["model1"]', "2025-01-01T00:00:00"))
|
||||||
|
conn.commit()
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
# Initialize with preserve flag
|
||||||
|
initialize_dev_database(db_path)
|
||||||
|
|
||||||
|
# Verify data is preserved
|
||||||
|
conn = get_db_connection(db_path)
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute("SELECT COUNT(*) FROM jobs")
|
||||||
|
assert cursor.fetchone()[0] == 1
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
os.environ.pop("PRESERVE_DEV_DATA")
|
||||||
Reference in New Issue
Block a user