Coverage for mlos_bench/mlos_bench/tests/storage/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 

7from alembic.migration import MigrationContext 

8from sqlalchemy import inspect 

9 

10from mlos_bench.storage.sql.storage import SqlStorage 

11 

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

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

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

15CURRENT_ALEMBIC_HEAD = "8928a401115b" 

16 

17 

18def test_storage_schemas(storage: SqlStorage) -> None: 

19 """Test storage schema creation.""" 

20 eng = storage._engine # pylint: disable=protected-access 

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

22 inspector = inspect(conn) 

23 # Make sure the "trial_runner_id" column exists. 

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

25 assert any( 

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

27 ) 

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

29 assert inspector.has_table("alembic_version") 

30 context = MigrationContext.configure(conn) 

31 current_rev = context.get_current_revision() 

32 assert ( 

33 current_rev == CURRENT_ALEMBIC_HEAD 

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