Skip to content

test: cover model-engine migration adoption #855

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ravneet-madaan-scale wants to merge 1 commit into from
Closed

test: cover model-engine migration adoption #855

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 169 additions & 0 deletions model-engine/tests/test_database_migration_adoption.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
from typing import Any
from unittest.mock import Mock, mock_open, patch

import pytest

VERSIONS_DIR = (
Path(__file__).parents[1] / "model_engine_server" / "db" / "migrations" / "alembic" / "versions"
)

INITIAL_MIGRATION = "2024_09_09_1736-fa3267c80731_initial.py"

ADD_COLUMN_MIGRATIONS = [
(
"2024_09_09_1831-b574e9711e35_chat_completion_add_extra_routes.py",
[("bundles", "runnable_image_extra_routes")],
),
(
"2024_09_24_1456-f55525c81eb5_multinode_bundle.py",
[
("bundles", "runnable_image_worker_command"),
("bundles", "runnable_image_worker_env"),
],
),
(
"2025_09_16_1741-e580182d6bfd_add_passthrough_forwarder.py",
[("bundles", "runnable_image_forwarder_type")],
),
(
"2025_09_25_1940-221aa19d3f32_add_routes_column.py",
[("bundles", "runnable_image_routes")],
),
(
"2026_02_10_1920-62da4f8b3403_add_task_expires_seconds_column.py",
[("endpoints", "task_expires_seconds")],
),
(
"2026_02_20_1200-a1b2c3d4e5f6_add_queue_message_timeout_seconds_column.py",
[("endpoints", "queue_message_timeout_seconds")],
),
(
"2026_06_16_1200-c4d5e6f7a8b9_add_status_reason_column.py",
[("endpoints", "status_reason")],
),
]


def load_migration(filename: str) -> Any:
path = VERSIONS_DIR / filename
spec = spec_from_file_location(path.stem.replace("-", "_"), path)
assert spec is not None
assert spec.loader is not None
module = module_from_spec(spec)
spec.loader.exec_module(module)
return module


def test_initial_migration_adopts_an_existing_schema(monkeypatch: pytest.MonkeyPatch) -> None:
migration = load_migration(INITIAL_MIGRATION)
inspector = Mock()
inspector.has_table.return_value = True
operation = Mock()
monkeypatch.setattr(migration.sa, "inspect", Mock(return_value=inspector))
monkeypatch.setattr(migration, "op", operation)

with patch("builtins.open", mock_open()) as open_file:
migration.upgrade()

inspector.has_table.assert_called_once_with("endpoints", schema="hosted_model_inference")
open_file.assert_not_called()
operation.execute.assert_not_called()


def test_initial_migration_initializes_an_empty_database(monkeypatch: pytest.MonkeyPatch) -> None:
migration = load_migration(INITIAL_MIGRATION)
inspector = Mock()
inspector.has_table.return_value = False
operation = Mock()
monkeypatch.setattr(migration.sa, "inspect", Mock(return_value=inspector))
monkeypatch.setattr(migration, "op", operation)

with patch("builtins.open", mock_open(read_data="SELECT 1;")) as open_file:
migration.upgrade()

open_file.assert_called_once_with(migration.INITIAL_MIGRATION_PATH)
operation.execute.assert_called_once_with("SELECT 1;")


@pytest.mark.parametrize(("filename", "columns"), ADD_COLUMN_MIGRATIONS)
def test_upgrade_skips_columns_that_already_exist(
filename: str,
columns: list[tuple[str, str]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
migration = load_migration(filename)
inspector = Mock()
inspector.get_columns.return_value = [{"name": column} for _, column in columns]
operation = Mock()
monkeypatch.setattr(migration.sa, "inspect", Mock(return_value=inspector))
monkeypatch.setattr(migration, "op", operation)

migration.upgrade()

operation.add_column.assert_not_called()
Comment on lines +88 to +105

@greptile-apps greptile-apps Bot Jul 15, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Schema argument not validated in get_columns mock

The inspector.get_columns mock always returns a fixed list regardless of the arguments it receives. None of the four test variants assert that get_columns was called with schema="hosted_model_inference". A migration that omits the schema keyword from its internal get_columns call would silently produce the same mock return value and the tests would still pass, masking the missing schema propagation. Adding inspector.get_columns.assert_any_call(table, schema="hosted_model_inference") (or equivalent per-call assertions) in the upgrade/downgrade tests would close this gap — the same applies to test_downgrade_skips_columns_that_are_already_absent and test_downgrade_drops_only_existing_columns.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: model-engine/tests/test_database_migration_adoption.py
Line: 88-105

Comment:
**Schema argument not validated in `get_columns` mock**

The `inspector.get_columns` mock always returns a fixed list regardless of the arguments it receives. None of the four test variants assert that `get_columns` was called with `schema="hosted_model_inference"`. A migration that omits the `schema` keyword from its internal `get_columns` call would silently produce the same mock return value and the tests would still pass, masking the missing schema propagation. Adding `inspector.get_columns.assert_any_call(table, schema="hosted_model_inference")` (or equivalent per-call assertions) in the upgrade/downgrade tests would close this gap — the same applies to `test_downgrade_skips_columns_that_are_already_absent` and `test_downgrade_drops_only_existing_columns`.

How can I resolve this? If you propose a fix, please make it concise.

Fix in Cursor Fix in Claude Code Fix in Codex



@pytest.mark.parametrize(("filename", "columns"), ADD_COLUMN_MIGRATIONS)
def test_upgrade_adds_only_missing_columns(
filename: str,
columns: list[tuple[str, str]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
migration = load_migration(filename)
inspector = Mock()
inspector.get_columns.return_value = []
operation = Mock()
monkeypatch.setattr(migration.sa, "inspect", Mock(return_value=inspector))
monkeypatch.setattr(migration, "op", operation)

migration.upgrade()

added_columns = [
(call.args[0], call.args[1].name, call.kwargs["schema"])
for call in operation.add_column.call_args_list
]
assert added_columns == [(table, column, "hosted_model_inference") for table, column in columns]


@pytest.mark.parametrize(("filename", "columns"), ADD_COLUMN_MIGRATIONS)
def test_downgrade_skips_columns_that_are_already_absent(
filename: str,
columns: list[tuple[str, str]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
migration = load_migration(filename)
inspector = Mock()
inspector.get_columns.return_value = []
operation = Mock()
monkeypatch.setattr(migration.sa, "inspect", Mock(return_value=inspector))
monkeypatch.setattr(migration, "op", operation)

migration.downgrade()

operation.drop_column.assert_not_called()


@pytest.mark.parametrize(("filename", "columns"), ADD_COLUMN_MIGRATIONS)
def test_downgrade_drops_only_existing_columns(
filename: str,
columns: list[tuple[str, str]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
migration = load_migration(filename)
inspector = Mock()
inspector.get_columns.return_value = [{"name": column} for _, column in columns]
operation = Mock()
monkeypatch.setattr(migration.sa, "inspect", Mock(return_value=inspector))
monkeypatch.setattr(migration, "op", operation)

migration.downgrade()

dropped_columns = [
(call.args[0], call.args[1], call.kwargs["schema"])
for call in operation.drop_column.call_args_list
]
assert dropped_columns == [
(table, column, "hosted_model_inference") for table, column in columns
]