103 lines
2.8 KiB
Python
103 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
from datetime import UTC
|
|
from pathlib import Path
|
|
|
|
from scripts.migrate_sqlite_to_postgres import (
|
|
REQUEST_COLUMNS,
|
|
USER_COLUMNS,
|
|
parse_bool,
|
|
parse_dt,
|
|
read_rows,
|
|
)
|
|
|
|
REQUEST_DEFAULTS = {
|
|
"username": None,
|
|
"enhanced_prompt": None,
|
|
"payment_charge_id": None,
|
|
"payment_refunded": 0,
|
|
"is_premium": 0,
|
|
"seed": None,
|
|
"error_text": None,
|
|
"finished_at": None,
|
|
}
|
|
|
|
USER_DEFAULTS = {
|
|
"username": None,
|
|
"credits": 3,
|
|
"last_daily_credit": "2026-10-18",
|
|
"referrer_id": None,
|
|
"first_paid_at": None,
|
|
"created_at": "2026-10-18T12:00:00+00:00",
|
|
"is_admin": 0,
|
|
"is_blocked": 0,
|
|
"age_confirmed": 0,
|
|
"age_declined": 0,
|
|
}
|
|
|
|
|
|
def test_read_rows_fills_legacy_defaults(tmp_path: Path) -> None:
|
|
path = tmp_path / "legacy.db"
|
|
connection = sqlite3.connect(path)
|
|
connection.row_factory = sqlite3.Row
|
|
connection.executescript(
|
|
"""
|
|
CREATE TABLE requests (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER NOT NULL,
|
|
timestamp TEXT NOT NULL,
|
|
prompt_text TEXT NOT NULL,
|
|
input_file_path TEXT,
|
|
output_file_path TEXT,
|
|
status TEXT NOT NULL DEFAULT 'queued',
|
|
comfy_prompt_id TEXT,
|
|
mode TEXT NOT NULL,
|
|
seed INTEGER,
|
|
error_text TEXT,
|
|
finished_at TEXT
|
|
);
|
|
CREATE TABLE users (
|
|
user_id INTEGER PRIMARY KEY,
|
|
username TEXT,
|
|
credits INTEGER NOT NULL DEFAULT 3,
|
|
last_daily_credit TEXT NOT NULL,
|
|
referrer_id INTEGER,
|
|
first_paid_at TEXT,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"INSERT INTO requests (user_id, timestamp, prompt_text, mode, seed)"
|
|
" VALUES (1, '2026-10-18T12:00:00+00:00', 'apple', 't2i', 42)"
|
|
)
|
|
connection.execute(
|
|
"INSERT INTO users"
|
|
" (user_id, username, credits, last_daily_credit, created_at)"
|
|
" VALUES (1, 'alice', 3, '2026-10-18',"
|
|
" '2026-10-18T12:00:00+00:00')"
|
|
)
|
|
connection.commit()
|
|
|
|
requests = read_rows(
|
|
connection, "requests", REQUEST_COLUMNS, REQUEST_DEFAULTS
|
|
)
|
|
users = read_rows(connection, "users", USER_COLUMNS, USER_DEFAULTS)
|
|
connection.close()
|
|
|
|
assert requests[0]["username"] is None
|
|
assert requests[0]["payment_refunded"] == 0
|
|
assert requests[0]["is_premium"] == 0
|
|
assert users[0]["is_admin"] == 0
|
|
assert users[0]["age_declined"] == 0
|
|
assert users[0]["username"] == "alice"
|
|
|
|
|
|
def test_parse_helpers() -> None:
|
|
parsed = parse_dt("2026-10-18T12:00:00")
|
|
assert parsed is not None
|
|
assert parsed.tzinfo is UTC
|
|
assert parse_bool("1") is True
|
|
assert parse_bool(0) is False
|