yuqei
86 lines
3.4 KiB
Python
86 lines
3.4 KiB
Python
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from sqlalchemy.sql.elements import TextClause
|
|
from sqlalchemy.exc import OperationalError
|
|
|
|
os.environ.setdefault("DATABASE_URL", "postgresql+psycopg://user:pass@localhost:5432/ai_ict_test")
|
|
|
|
from scripts import run_migrations as run_migrations_module
|
|
from scripts.run_migrations import apply_migration_file, run_migrations as run_migration_files, split_sql_statements
|
|
|
|
|
|
class MigrationRunnerTests(unittest.TestCase):
|
|
def test_split_sql_statements_keeps_semicolon_inside_quotes(self) -> None:
|
|
sql = """
|
|
-- comment with ; should be ignored
|
|
create table sample (id bigint);
|
|
insert into sample (note) values ('semi;colon');
|
|
"""
|
|
|
|
statements = split_sql_statements(sql)
|
|
|
|
self.assertEqual(len(statements), 2)
|
|
self.assertIn("create table sample", statements[0])
|
|
self.assertIn("'semi;colon'", statements[1])
|
|
|
|
def test_apply_migration_file_executes_each_statement(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
path = Path(tmpdir) / "001_sample.sql"
|
|
path.write_text("create table a (id bigint); create table b (id bigint);", encoding="utf-8")
|
|
connection = FakeConnection()
|
|
|
|
count = apply_migration_file(connection, path)
|
|
|
|
self.assertEqual(count, 2)
|
|
self.assertEqual(len(connection.statements), 2)
|
|
self.assertTrue(all(isinstance(statement, TextClause) for statement in connection.statements))
|
|
|
|
def test_run_migrations_reports_each_file(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
first = Path(tmpdir) / "001_sample.sql"
|
|
second = Path(tmpdir) / "002_sample.sql"
|
|
first.write_text("create table a (id bigint);", encoding="utf-8")
|
|
second.write_text("create table b (id bigint); create table c (id bigint);", encoding="utf-8")
|
|
connection = FakeConnection()
|
|
|
|
results = run_migration_files(connection=connection, files=[first, second])
|
|
|
|
self.assertEqual(results, [{"file": "001_sample.sql", "statements": 1}, {"file": "002_sample.sql", "statements": 2}])
|
|
|
|
def test_main_reports_database_unavailable(self) -> None:
|
|
output = io.StringIO()
|
|
|
|
with patch("scripts.run_migrations.engine.begin", side_effect=OperationalError("stmt", {}, Exception())), contextlib.redirect_stdout(output):
|
|
exit_code = run_migrations_module.main()
|
|
|
|
self.assertEqual(exit_code, 2)
|
|
self.assertIn("Database unavailable for migrations: OperationalError", output.getvalue())
|
|
|
|
def test_main_can_emit_json_failure(self) -> None:
|
|
output = io.StringIO()
|
|
|
|
with patch("scripts.run_migrations.engine.begin", side_effect=OperationalError("stmt", {}, Exception())), contextlib.redirect_stdout(output):
|
|
exit_code = run_migrations_module.main(env={"MIGRATIONS_OUTPUT_JSON": "1"})
|
|
|
|
payload = json.loads(output.getvalue())
|
|
self.assertEqual(exit_code, 2)
|
|
self.assertFalse(payload["passed"])
|
|
self.assertEqual(payload["error_type"], "OperationalError")
|
|
|
|
|
|
class FakeConnection:
|
|
def __init__(self) -> None:
|
|
self.statements = []
|
|
|
|
def execute(self, statement):
|
|
self.statements.append(statement)
|
|
if __name__ == "__main__":
|
|
unittest.main()
|