refactoring storage tests to check other db engines

microsoft · bpkroth · May 28, 2025 · May 28, 2025 · May 28, 2025 · May 28, 2025
commit 3b3bf6ce50ff23e43bd0bf62c26f2410877ea458
@@ -296,6 +296,31 @@ def _get_alembic_cfg(conn: Connection) -> config.Config:
         alembic_cfg.attributes["connection"] = conn
         return alembic_cfg
 
+    def drop_all_tables(self, *, force: bool = False) -> None:
+        """
+        Helper method used in testing to reset the DB schema.
+
+        Notes
+        -----
+        This method is not intended for production use, as it will drop all tables
+        in the database. Use with caution.
+
+        Parameters
+        ----------
+        force : bool
+            If True, drop all tables in the target database.
+            If False, this method will not drop any tables and will log a warning.
+        """
+        assert self._engine
+        self.meta.reflect(bind=self._engine)
+        if force:
+            self.meta.drop_all(bind=self._engine)
+        else:
+            _LOG.warning(
+                "Resetting the schema without force is not implemented. "
+                "Use force=True to drop all tables."
+            )
+
     def create(self) -> "DbSchema":
         """Create the DB schema."""
         _LOG.info("Create the DB schema")

@@ -82,6 +82,33 @@ def _schema(self) -> DbSchema:
                 _LOG.debug("DDL statements:\n%s", self._db_schema)
         return self._db_schema
 
+    def _reset_schema(self, *, force: bool = False) -> None:
+        """
+        Helper method used in testing to reset the DB schema.
+
+        Notes
+        -----
+        This method is not intended for production use, as it will drop all tables
+        in the database. Use with caution.
+
+        Parameters
+        ----------
+        force : bool
+            If True, drop all tables in the target database.
+            If False, this method will not drop any tables and will log a warning.
+        """
+        assert self._engine
+        if force:
+            self._schema.drop_all_tables(force=force)
+            self._db_schema = DbSchema(self._engine)
+            self._schema_created = False
+            self._schema_updated = False
+        else:
+            _LOG.warning(
+                "Resetting the schema without force is not implemented. "
+                "Use force=True to drop all tables."
+            )
+
     def update_schema(self) -> None:
         """Update the database schema."""
         if not self._schema_updated:

@@ -86,6 +86,30 @@ def check_class_name(obj: object, expected_class_name: str) -> bool:
     full_class_name = obj.__class__.__module__ + "." + obj.__class__.__name__
     return full_class_name == try_resolve_class_name(expected_class_name)
 
+HOST_DOCKER_NAME = "host.docker.internal"
+
+
+@pytest.fixture(scope="session")
+def docker_hostname() -> str:
+    """Returns the local hostname to use to connect to the test ssh server."""
+    if sys.platform != "win32" and resolve_host_name(HOST_DOCKER_NAME):
+        # On Linux, if we're running in a docker container, we can use the
+        # --add-host (extra_hosts in docker-compose.yml) to refer to the host IP.
+        return HOST_DOCKER_NAME
+    # Docker (Desktop) for Windows (WSL2) uses a special networking magic
+    # to refer to the host machine as `localhost` when exposing ports.
+    # In all other cases, assume we're executing directly inside conda on the host.
+    return "localhost"
+
+
+def wait_docker_service_socket(docker_services: DockerServices, hostname: str, port: int) -> None:
+    """Wait until a docker service is ready."""
+    docker_services.wait_until_responsive(
+        check=lambda: check_socket(hostname, port),
+        timeout=30.0,
+        pause=0.5,
+    )
+
 
 def check_socket(host: str, port: int, timeout: float = 1.0) -> bool:
     """

@@ -14,7 +14,7 @@
 from mlos_bench.schedulers.base_scheduler import Scheduler
 from mlos_bench.schedulers.trial_runner import TrialRunner
 from mlos_bench.services.config_persistence import ConfigPersistenceService
-from mlos_bench.storage.sql.storage import SqlStorage
+from mlos_bench.storage.base_storage import Storage
 from mlos_bench.tests.config import BUILTIN_TEST_CONFIG_PATH, locate_config_examples
 from mlos_bench.util import get_class_from_name
 
@@ -58,7 +58,7 @@ def test_load_scheduler_config_examples(
     config_path: str,
     mock_env_config_path: str,
     trial_runners: list[TrialRunner],
-    storage: SqlStorage,
+    storage: Storage,
     mock_opt: MockOptimizer,
 ) -> None:
     """Tests loading a config example."""

@@ -90,6 +90,7 @@ def docker_compose_file(pytestconfig: pytest.Config) -> list[str]:
     _ = pytestconfig  # unused
     return [
         os.path.join(os.path.dirname(__file__), "services", "remote", "ssh", "docker-compose.yml"),
+        os.path.join(os.path.dirname(__file__), "storage", "sql", "docker-compose.yml"),
         # Add additional configs as necessary here.
     ]
 

@@ -22,7 +22,7 @@ These are brought up as session fixtures under a unique (PID based) compose proj
 
 In the case of `pytest`, since the `SshService` base class implements a shared connection cache that we wish to test, and testing "rebooting" of servers (containers) is also necessary, but we want to avoid single threaded execution for tests, we start a third container only for testing reboots.
 
-Additionally, since `scope="session"` fixtures are executed once per worker, which is excessive in our case, we use lockfiles (one of the only portal synchronization methods) to ensure that the usual `docker_services` fixture which starts and stops the containers is only executed once per test run and uses a shared compose instance.
+Additionally, since `scope="session"` fixtures are executed once per worker, which is excessive in our case, we use lockfiles (one of the only portable synchronization methods) to ensure that the usual `docker_services` fixture which starts and stops the containers is only executed once per test run and uses a shared compose instance.
 
 ## See Also
 

@@ -7,9 +7,6 @@
 from dataclasses import dataclass
 from subprocess import run
 
-from pytest_docker.plugin import Services as DockerServices
-
-from mlos_bench.tests import check_socket
 
 # The SSH test server port and name.
 # See Also: docker-compose.yml
@@ -21,7 +18,12 @@
 
 @dataclass
 class SshTestServerInfo:
-    """A data class for SshTestServerInfo."""
+    """A data class for SshTestServerInfo.
+
+    See Also
+    --------
+    mlos_bench.tests.storage.sql.SqlTestServerInfo
+    """
 
     compose_project_name: str
     service_name: str
@@ -70,12 +72,3 @@ def to_connect_params(self, uncached: bool = False) -> dict:
             "port": self.get_port(uncached),
             "username": self.username,
         }
-
-
-def wait_docker_service_socket(docker_services: DockerServices, hostname: str, port: int) -> None:
-    """Wait until a docker service is ready."""
-    docker_services.wait_until_responsive(
-        check=lambda: check_socket(hostname, port),
-        timeout=30.0,
-        pause=0.5,
-    )
@@ -19,36 +19,20 @@
 
 from mlos_bench.services.remote.ssh.ssh_fileshare import SshFileShareService
 from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService
-from mlos_bench.tests import resolve_host_name
+from mlos_bench.tests import wait_docker_service_socket
 from mlos_bench.tests.services.remote.ssh import (
     ALT_TEST_SERVER_NAME,
     REBOOT_TEST_SERVER_NAME,
     SSH_TEST_SERVER_NAME,
     SshTestServerInfo,
-    wait_docker_service_socket,
 )
 
 # pylint: disable=redefined-outer-name
 
-HOST_DOCKER_NAME = "host.docker.internal"
-
-
-@pytest.fixture(scope="session")
-def ssh_test_server_hostname() -> str:
-    """Returns the local hostname to use to connect to the test ssh server."""
-    if sys.platform != "win32" and resolve_host_name(HOST_DOCKER_NAME):
-        # On Linux, if we're running in a docker container, we can use the
-        # --add-host (extra_hosts in docker-compose.yml) to refer to the host IP.
-        return HOST_DOCKER_NAME
-    # Docker (Desktop) for Windows (WSL2) uses a special networking magic
-    # to refer to the host machine as `localhost` when exposing ports.
-    # In all other cases, assume we're executing directly inside conda on the host.
-    return "localhost"
-
 
 @pytest.fixture(scope="session")
 def ssh_test_server(
-    ssh_test_server_hostname: str,
+    docker_hostname: str,
     docker_compose_project_name: str,
     locked_docker_services: DockerServices,
 ) -> Generator[SshTestServerInfo]:
@@ -66,12 +50,14 @@ def ssh_test_server(
         ssh_test_server_info = SshTestServerInfo(
             compose_project_name=docker_compose_project_name,
             service_name=SSH_TEST_SERVER_NAME,
-            hostname=ssh_test_server_hostname,
+            hostname=docker_hostname,
             username="root",
             id_rsa_path=id_rsa_file.name,
         )
         wait_docker_service_socket(
-            locked_docker_services, ssh_test_server_info.hostname, ssh_test_server_info.get_port()
+            locked_docker_services,
+            ssh_test_server_info.hostname,
+            ssh_test_server_info.get_port(),
         )
         id_rsa_src = f"/{ssh_test_server_info.username}/.ssh/id_rsa"
         docker_cp_cmd = (
@@ -116,7 +102,9 @@ def alt_test_server(
         id_rsa_path=ssh_test_server.id_rsa_path,
     )
     wait_docker_service_socket(
-        locked_docker_services, alt_test_server_info.hostname, alt_test_server_info.get_port()
+        locked_docker_services,
+        alt_test_server_info.hostname,
+        alt_test_server_info.get_port(),
     )
     return alt_test_server_info
 

@@ -12,13 +12,12 @@
 
 from mlos_bench.services.remote.ssh.ssh_host_service import SshHostService
 from mlos_bench.services.remote.ssh.ssh_service import SshClient
-from mlos_bench.tests import requires_docker
+from mlos_bench.tests import requires_docker, wait_docker_service_socket,
 from mlos_bench.tests.services.remote.ssh import (
     ALT_TEST_SERVER_NAME,
     REBOOT_TEST_SERVER_NAME,
     SSH_TEST_SERVER_NAME,
     SshTestServerInfo,
-    wait_docker_service_socket,
 )
 
 _LOG = logging.getLogger(__name__)

@@ -11,8 +11,10 @@
 # same.
 
 # Expose some of those as local names so they can be picked up as fixtures by pytest.
-storage = sql_storage_fixtures.storage
+mysql_storage = sql_storage_fixtures.mysql_storage
+postgres_storage = sql_storage_fixtures.postgres_storage
 sqlite_storage = sql_storage_fixtures.sqlite_storage
+storage = sql_storage_fixtures.storage
 exp_storage = sql_storage_fixtures.exp_storage
 exp_no_tunables_storage = sql_storage_fixtures.exp_no_tunables_storage
 mixed_numerics_exp_storage = sql_storage_fixtures.mixed_numerics_exp_storage

@@ -0,0 +1,27 @@
+# Sql Storage Tests
+
+The "unit" tests for the `SqlStorage` classes are more functional than other unit tests in that we don't merely mock them out, but actually setup small SQL databases with `docker compose` and interact with them using the `SqlStorage` class.
+
+To do this, we make use of the `pytest-docker` plugin to bring up the services defined in the [`docker-compose.yml`](./docker-compose.yml) file in this directory.
+
+There are currently two services defined in that config, though others could be added in the future:
+
+1. `mysql-mlos-bench-server`
+1. `postgres-mlos-bench-server`
+
+We rely on `docker compose` to map their internal container service ports to random ports on the host.
+Hence, when connecting, we need to look up these ports on demand using something akin to `docker compose port`.
+Because of complexities of networking in different development environments (especially for Docker on WSL2 for Windows), we may also have to connect to a different host address than `localhost` (e.g., `host.docker.internal`, which is dynamically requested as a part of of the [devcontainer](../../../../../../.devcontainer/docker-compose.yml) setup).
+
+These containers are brought up as session fixtures under a unique (PID based) compose project name for each `pytest` invocation, but only when docker is detected on the host (via the `@docker_required` decorator we define in [`mlos_bench/tests/__init__.py`](../../../__init__.py)), else those tests are skipped.
+
+> For manual testing, to bring up/down the test infrastructure the [`up.sh`](./up.sh) and [`down.sh`](./down.sh) scripts can be used, which assigns a known project name.
+
+In the case of `pytest`, we also want to be able to test with a fresh state in most cases, so we use the `pytest` `yield` pattern to allow schema cleanup code to happen after the end of each test (see: `_create_storage_from_test_server_info`).
+We use lockfiles to prevent races between tests that would otherwise try to create or cleanup the same database schema at the same time.
+
+Additionally, since `scope="session"` fixtures are executed once per worker, which is excessive in our case, we use lockfiles (one of the only portable synchronization methods) to ensure that the usual `docker_services` fixture which starts and stops the containers is only executed once per test run and uses a shared compose instance.
+
+## See Also
+
+Notes in the [`mlos_bench/tests/services/remote/ssh/README.md`](../../../services/remote/ssh/README.md) file for a similar setup for SSH services.
@@ -3,3 +3,75 @@
 # Licensed under the MIT License.
 #
 """Tests for mlos_bench sql storage."""
+
+from dataclasses import dataclass
+from subprocess import run
+
+
+# The DB servers' names and other connection info.
+# See Also: docker-compose.yml
+
+MYSQL_TEST_SERVER_NAME = "mysql-mlos-bench-server"
+PGSQL_TEST_SERVER_NAME = "postgres-mlos-bench-server"
+
+SQL_TEST_SERVER_DATABASE = "mlos_bench"
+SQL_TEST_SERVER_PASSWORD = "password"
+
+
+@dataclass
+class SqlTestServerInfo:
+    """A data class for SqlTestServerInfo.
+
+    See Also
+    --------
+    mlos_bench.tests.services.remote.ssh.SshTestServerInfo
+    """
+
+    compose_project_name: str
+    service_name: str
+    hostname: str
+    _port: int | None = None
+
+    @property
+    def username(self) -> str:
+        """Gets the username."""
+        usernames = {
+            MYSQL_TEST_SERVER_NAME: "root",
+            PGSQL_TEST_SERVER_NAME: "postgres",
+        }
+        return usernames[self.service_name]
+
+    @property
+    def password(self) -> str:
+        """Gets the password."""
+        return SQL_TEST_SERVER_PASSWORD
+
+    @property
+    def database(self) -> str:
+        """Gets the database."""
+        return SQL_TEST_SERVER_DATABASE
+
+    def get_port(self, uncached: bool = False) -> int:
+        """
+        Gets the port that the SSH test server is listening on.
+
+        Note: this value can change when the service restarts so we can't rely on
+        the DockerServices.
+        """
+        if self._port is None or uncached:
+            default_ports = {
+                MYSQL_TEST_SERVER_NAME: 3306,
+                PGSQL_TEST_SERVER_NAME: 5432,
+            }
+            default_port = default_ports[self.service_name]
+            port_cmd = run(
+                (
+                    f"docker compose -p {self.compose_project_name} "
+                    f"port {self.service_name} {default_port}"
+                ),
+                shell=True,
+                check=True,
+                capture_output=True,
+            )
+            self._port = int(port_cmd.stdout.decode().strip().split(":")[1])
+        return self._port