test(cloud-store): verify schema migrations
This commit is contained in:
@@ -0,0 +1,140 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import create_engine, inspect, text
|
||||
|
||||
from cloud.schema import (
|
||||
HEAD_REVISION,
|
||||
SchemaVersionError,
|
||||
current_revision,
|
||||
downgrade_database,
|
||||
require_current_schema,
|
||||
upgrade_database,
|
||||
)
|
||||
|
||||
|
||||
def _database_url(tmp_path) -> str:
|
||||
return f"sqlite:///{(tmp_path / 'cloud.sqlite3').as_posix()}"
|
||||
|
||||
|
||||
def test_forward_and_downgrade_migrations_on_empty_database(tmp_path) -> None:
|
||||
database_url = _database_url(tmp_path)
|
||||
|
||||
upgrade_database(database_url)
|
||||
|
||||
engine = create_engine(database_url)
|
||||
try:
|
||||
table_names = set(inspect(engine).get_table_names())
|
||||
assert {
|
||||
"host_registrations",
|
||||
"pooled_devices",
|
||||
"scheduled_tasks",
|
||||
"plugins",
|
||||
"task_attempts",
|
||||
} <= table_names
|
||||
assert current_revision(database_url) == HEAD_REVISION
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
downgrade_database(database_url)
|
||||
|
||||
engine = create_engine(database_url)
|
||||
try:
|
||||
inspector = inspect(engine)
|
||||
assert "task_attempts" not in inspector.get_table_names()
|
||||
task_columns = {
|
||||
column["name"] for column in inspector.get_columns("scheduled_tasks")
|
||||
}
|
||||
assert "lease_id" not in task_columns
|
||||
assert current_revision(database_url) is None
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
|
||||
def test_legacy_data_survives_upgrade_and_downgrade(tmp_path) -> None:
|
||||
database_url = _database_url(tmp_path)
|
||||
engine = create_engine(database_url)
|
||||
try:
|
||||
with engine.begin() as connection:
|
||||
_create_legacy_schema(connection)
|
||||
connection.execute(
|
||||
text(
|
||||
"insert into host_registrations "
|
||||
"(host_id, address, last_seen_at) values "
|
||||
"('host-a', 'local', '2026-01-01T00:00:00+00:00')"
|
||||
)
|
||||
)
|
||||
connection.execute(
|
||||
text(
|
||||
"insert into scheduled_tasks "
|
||||
"(id, goal, workflow_definition_id, constraints_json, status, "
|
||||
"assigned_device_id, assigned_host_id, created_at) values "
|
||||
"('task-a', 'goal', null, :constraints, 'queued', null, null, "
|
||||
"'2026-01-01T00:00:00+00:00')"
|
||||
),
|
||||
{"constraints": json.dumps({"driver_type": None, "capability_tags": []})},
|
||||
)
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
upgrade_database(database_url)
|
||||
engine = create_engine(database_url)
|
||||
try:
|
||||
with engine.connect() as connection:
|
||||
assert connection.scalar(text("select count(*) from host_registrations")) == 1
|
||||
assert connection.scalar(text("select count(*) from scheduled_tasks")) == 1
|
||||
task_columns = {
|
||||
column["name"] for column in inspect(engine).get_columns("scheduled_tasks")
|
||||
}
|
||||
assert {"attempt_count", "lease_id", "result_json"} <= task_columns
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
downgrade_database(database_url)
|
||||
engine = create_engine(database_url)
|
||||
try:
|
||||
with engine.connect() as connection:
|
||||
assert connection.scalar(text("select count(*) from host_registrations")) == 1
|
||||
assert connection.scalar(text("select count(*) from scheduled_tasks")) == 1
|
||||
task_columns = {
|
||||
column["name"] for column in inspect(engine).get_columns("scheduled_tasks")
|
||||
}
|
||||
assert "attempt_count" not in task_columns
|
||||
finally:
|
||||
engine.dispose()
|
||||
|
||||
|
||||
def test_schema_readiness_requires_head_revision(tmp_path) -> None:
|
||||
database_url = _database_url(tmp_path)
|
||||
|
||||
with pytest.raises(SchemaVersionError, match="unversioned"):
|
||||
require_current_schema(database_url)
|
||||
|
||||
upgrade_database(database_url)
|
||||
require_current_schema(database_url)
|
||||
|
||||
|
||||
def _create_legacy_schema(connection) -> None:
|
||||
connection.exec_driver_sql(
|
||||
"create table host_registrations ("
|
||||
"host_id text primary key, address text, last_seen_at text not null)"
|
||||
)
|
||||
connection.exec_driver_sql(
|
||||
"create table pooled_devices ("
|
||||
"device_id text not null, host_id text not null, driver_type text not null, "
|
||||
"status text not null, capability_tags_json text not null, synced_at text, "
|
||||
"primary key (host_id, device_id))"
|
||||
)
|
||||
connection.exec_driver_sql(
|
||||
"create table scheduled_tasks ("
|
||||
"id text primary key, goal text, workflow_definition_id text, "
|
||||
"constraints_json text not null, status text not null, "
|
||||
"assigned_device_id text, assigned_host_id text, created_at text not null)"
|
||||
)
|
||||
connection.exec_driver_sql(
|
||||
"create table plugins ("
|
||||
"name text primary key, version text not null, entry_point_kind text not null, "
|
||||
"target text not null, wired integer not null)"
|
||||
)
|
||||
Reference in New Issue
Block a user