Coverage for mlos_bench/mlos_bench/tests/storage/sql/test_storage_schemas.py: 100%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Test sql schemas for mlos_bench storage.""" 

6 

7import pytest 

8from alembic.migration import MigrationContext 

9from pytest_lazy_fixtures.lazy_fixture import lf as lazy_fixture 

10from sqlalchemy import inspect 

11 

12from mlos_bench.storage.sql.storage import SqlStorage 

13from mlos_bench.tests.storage.sql.fixtures import DOCKER_DBMS_FIXTURES 

14 

15# NOTE: This value is hardcoded to the latest revision in the alembic versions directory. 

16# It could also be obtained programmatically using the "alembic heads" command or heads() API. 

17# See Also: schema.py for an example of programmatic alembic config access. 

18CURRENT_ALEMBIC_HEAD = "b61aa446e724" 

19 

20# Try to test multiple DBMS engines. 

21 

22 

23@pytest.mark.parametrize( 

24 "some_sql_storage_fixture", 

25 [ 

26 lazy_fixture("sqlite_storage"), 

27 *DOCKER_DBMS_FIXTURES, 

28 ], 

29) 

30def test_storage_schemas(some_sql_storage_fixture: SqlStorage) -> None: 

31 """Test storage schema creation.""" 

32 assert isinstance(some_sql_storage_fixture, SqlStorage) 

33 eng = some_sql_storage_fixture._engine # pylint: disable=protected-access 

34 with eng.connect() as conn: # pylint: disable=protected-access 

35 inspector = inspect(conn) 

36 # Make sure the "trial_runner_id" column exists. 

37 # (i.e., the latest schema has been applied) 

38 assert any( 

39 column["name"] == "trial_runner_id" for column in inspect(conn).get_columns("trial") 

40 ) 

41 # Make sure the "alembic_version" table exists and is appropriately stamped. 

42 assert inspector.has_table("alembic_version") 

43 context = MigrationContext.configure(conn) 

44 current_rev = context.get_current_revision() 

45 assert ( 

46 current_rev == CURRENT_ALEMBIC_HEAD 

47 ), f"Expected {CURRENT_ALEMBIC_HEAD}, got {current_rev}"