|
| 1 | +#!/usr/bin/env python |
| 2 | +"""Pytest configuration for singlestoredb tests |
| 3 | +
|
| 4 | +This module sets up automatic Docker container management for tests. |
| 5 | +It works with both pytest-style and unittest-style tests. |
| 6 | +
|
| 7 | +The conftest automatically: |
| 8 | +1. Checks if SINGLESTOREDB_URL is set in the environment |
| 9 | +2. If not set, starts a SingleStore Docker container |
| 10 | +3. Sets SINGLESTOREDB_URL for all tests to use |
| 11 | +4. Cleans up the container when tests complete |
| 12 | +
|
| 13 | +Environment Variables: |
| 14 | + - SINGLESTOREDB_URL: If set, tests will use this existing server instead |
| 15 | + of starting a Docker container. This allows testing |
| 16 | + against a specific server instance. |
| 17 | + - USE_DATA_API: If set to 1/true/on, tests will use HTTP Data API |
| 18 | + instead of MySQL protocol. When set, SINGLESTOREDB_URL |
| 19 | + will be set to the HTTP URL, and SINGLESTOREDB_INIT_DB_URL |
| 20 | + will be set to the MySQL URL for setup operations. |
| 21 | + - SINGLESTORE_LICENSE: Optional. License key for Docker container. If not |
| 22 | + set, an empty string is used as fallback. |
| 23 | +
|
| 24 | +Available Fixtures: |
| 25 | + - singlestoredb_test_container: Manages Docker container lifecycle |
| 26 | + - singlestoredb_connection: Provides a connection to the test server |
| 27 | + - singlestoredb_tempdb: Creates a temporary test database with cursor |
| 28 | +""" |
| 29 | +import logging |
| 30 | +import os |
| 31 | +from collections.abc import Iterator |
| 32 | +from typing import Optional |
| 33 | + |
| 34 | +import pytest |
| 35 | + |
| 36 | +from singlestoredb.pytest import _TestContainerManager |
| 37 | +from singlestoredb.pytest import execution_mode # noqa: F401 |
| 38 | +from singlestoredb.pytest import ExecutionMode # noqa: F401 |
| 39 | +from singlestoredb.pytest import name_allocator # noqa: F401 |
| 40 | +from singlestoredb.pytest import node_name # noqa: F401 |
| 41 | +from singlestoredb.pytest import singlestoredb_connection # noqa: F401 |
| 42 | +from singlestoredb.pytest import singlestoredb_tempdb # noqa: F401 |
| 43 | + |
| 44 | + |
| 45 | +logger = logging.getLogger(__name__) |
| 46 | + |
| 47 | +# Global container manager instance |
| 48 | +_container_manager: Optional[_TestContainerManager] = None |
| 49 | + |
| 50 | + |
| 51 | +def pytest_configure(config: pytest.Config) -> None: |
| 52 | + """ |
| 53 | + Pytest hook that runs before test collection. |
| 54 | +
|
| 55 | + This ensures the Docker container is started (if needed) before any |
| 56 | + test modules are imported. Some test modules try to get connection |
| 57 | + parameters at import time, so we need the environment set up early. |
| 58 | + """ |
| 59 | + global _container_manager |
| 60 | + |
| 61 | + # Prevent double initialization - pytest_configure can be called multiple times |
| 62 | + if _container_manager is not None: |
| 63 | + logger.debug('pytest_configure already called, skipping') |
| 64 | + return |
| 65 | + |
| 66 | + if 'SINGLESTOREDB_URL' not in os.environ: |
| 67 | + print('\n' + '=' * 70) |
| 68 | + print('Starting SingleStoreDB Docker container...') |
| 69 | + print('This may take a moment...') |
| 70 | + print('=' * 70) |
| 71 | + logger.info('SINGLESTOREDB_URL not set, starting Docker container') |
| 72 | + |
| 73 | + # Create and start the container |
| 74 | + _container_manager = _TestContainerManager() |
| 75 | + |
| 76 | + if not _container_manager.use_existing: |
| 77 | + _container_manager.start() |
| 78 | + print(f'Container {_container_manager.container_name} started') |
| 79 | + print('Waiting for SingleStoreDB to be ready...') |
| 80 | + |
| 81 | + # Wait for container to be ready |
| 82 | + try: |
| 83 | + conn = _container_manager.connect() |
| 84 | + conn.close() |
| 85 | + print('✓ SingleStoreDB is ready!') |
| 86 | + logger.info('Docker container is ready') |
| 87 | + except Exception as e: |
| 88 | + print(f'✗ Failed to connect to Docker container: {e}') |
| 89 | + logger.error(f'Failed to connect to Docker container: {e}') |
| 90 | + raise |
| 91 | + |
| 92 | + # Set the environment variable for all tests |
| 93 | + # Check if USE_DATA_API is set to use HTTP connection |
| 94 | + if os.environ.get('USE_DATA_API', '0').lower() in ('1', 'true', 'on'): |
| 95 | + # Use HTTP URL for tests |
| 96 | + url = _container_manager.http_connection_url |
| 97 | + if url is None: |
| 98 | + raise RuntimeError( |
| 99 | + 'Failed to get HTTP URL from container manager', |
| 100 | + ) |
| 101 | + os.environ['SINGLESTOREDB_URL'] = url |
| 102 | + print('=' * 70) |
| 103 | + print('USE_DATA_API is enabled - using HTTP Data API for tests') |
| 104 | + print(f'Tests will connect via: {url}') |
| 105 | + print('=' * 70) |
| 106 | + logger.info('USE_DATA_API is enabled - using HTTP Data API for tests') |
| 107 | + logger.info(f'Tests will connect via: {url}') |
| 108 | + |
| 109 | + # Also set INIT_DB_URL to MySQL URL for setup operations |
| 110 | + # (like SET GLOBAL) that don't work over HTTP |
| 111 | + mysql_url = _container_manager.url |
| 112 | + if mysql_url is None: |
| 113 | + raise RuntimeError( |
| 114 | + 'Failed to get MySQL URL from container manager', |
| 115 | + ) |
| 116 | + os.environ['SINGLESTOREDB_INIT_DB_URL'] = mysql_url |
| 117 | + print(f'Setup operations will use MySQL protocol: {mysql_url}') |
| 118 | + logger.info( |
| 119 | + f'Setup operations will use MySQL protocol: {mysql_url}', |
| 120 | + ) |
| 121 | + else: |
| 122 | + url = _container_manager.url |
| 123 | + if url is None: |
| 124 | + raise RuntimeError( |
| 125 | + 'Failed to get database URL from container manager', |
| 126 | + ) |
| 127 | + os.environ['SINGLESTOREDB_URL'] = url |
| 128 | + print('=' * 70) |
| 129 | + print(f'Tests will connect via MySQL protocol: {url}') |
| 130 | + print('=' * 70) |
| 131 | + logger.info(f'Tests will connect via MySQL protocol: {url}') |
| 132 | + else: |
| 133 | + url = os.environ['SINGLESTOREDB_URL'] |
| 134 | + logger.debug(f'Using existing SINGLESTOREDB_URL={url}') |
| 135 | + |
| 136 | + |
| 137 | +def pytest_unconfigure(config: pytest.Config) -> None: |
| 138 | + """ |
| 139 | + Pytest hook that runs after all tests complete. |
| 140 | +
|
| 141 | + Cleans up the Docker container if one was started. |
| 142 | + """ |
| 143 | + global _container_manager |
| 144 | + |
| 145 | + if _container_manager is not None and not _container_manager.use_existing: |
| 146 | + print('\n' + '=' * 70) |
| 147 | + print('Cleaning up Docker container...') |
| 148 | + logger.info('Cleaning up Docker container') |
| 149 | + try: |
| 150 | + _container_manager.stop() |
| 151 | + print(f'✓ Container {_container_manager.container_name} stopped') |
| 152 | + print('=' * 70) |
| 153 | + logger.info('Docker container stopped') |
| 154 | + except Exception as e: |
| 155 | + print(f'✗ Failed to stop Docker container: {e}') |
| 156 | + print('=' * 70) |
| 157 | + logger.error(f'Failed to stop Docker container: {e}') |
| 158 | + |
| 159 | + |
| 160 | +@pytest.fixture(scope='session', autouse=True) |
| 161 | +def setup_test_environment() -> Iterator[None]: |
| 162 | + """ |
| 163 | + Automatically set up test environment for all tests. |
| 164 | +
|
| 165 | + This fixture ensures the test environment is ready. The actual container |
| 166 | + setup happens in pytest_configure hook to ensure it runs before test |
| 167 | + collection. Cleanup happens in pytest_unconfigure hook. |
| 168 | +
|
| 169 | + This fixture exists to ensure proper ordering but doesn't manage the |
| 170 | + container lifecycle itself. |
| 171 | + """ |
| 172 | + # The environment should already be set up by pytest_configure |
| 173 | + # This fixture just ensures proper test initialization order |
| 174 | + yield |
| 175 | + |
| 176 | + # Clean up is handled by pytest_unconfigure |
| 177 | + |
| 178 | + |
| 179 | +@pytest.fixture(autouse=True) |
| 180 | +def protect_singlestoredb_url() -> Iterator[None]: |
| 181 | + """ |
| 182 | + Protect SINGLESTOREDB_URL and SINGLESTOREDB_INIT_DB_URL from corruption. |
| 183 | +
|
| 184 | + Some tests (like test_config.py) call reset_option() which resets all |
| 185 | + config options to their defaults. Since the 'host' option is registered |
| 186 | + with environ=['SINGLESTOREDB_HOST', 'SINGLESTOREDB_URL'], resetting it |
| 187 | + overwrites SINGLESTOREDB_URL with just '127.0.0.1' instead of the full |
| 188 | + connection string, breaking subsequent tests. |
| 189 | +
|
| 190 | + This fixture saves both URLs before each test and restores them |
| 191 | + after, ensuring they're not corrupted. |
| 192 | + """ |
| 193 | + # Save the current URLs |
| 194 | + saved_url = os.environ.get('SINGLESTOREDB_URL') |
| 195 | + saved_init_url = os.environ.get('SINGLESTOREDB_INIT_DB_URL') |
| 196 | + |
| 197 | + yield |
| 198 | + |
| 199 | + # Restore SINGLESTOREDB_URL if it was set and has been corrupted |
| 200 | + if saved_url is not None: |
| 201 | + current_url = os.environ.get('SINGLESTOREDB_URL') |
| 202 | + if current_url != saved_url: |
| 203 | + logger.debug( |
| 204 | + f'Restoring SINGLESTOREDB_URL from {current_url!r} to {saved_url!r}', |
| 205 | + ) |
| 206 | + os.environ['SINGLESTOREDB_URL'] = saved_url |
| 207 | + |
| 208 | + # Restore SINGLESTOREDB_INIT_DB_URL if it was set and has been corrupted |
| 209 | + if saved_init_url is not None: |
| 210 | + current_init_url = os.environ.get('SINGLESTOREDB_INIT_DB_URL') |
| 211 | + if current_init_url != saved_init_url: |
| 212 | + logger.debug( |
| 213 | + f'Restoring SINGLESTOREDB_INIT_DB_URL from ' |
| 214 | + f'{current_init_url!r} to {saved_init_url!r}', |
| 215 | + ) |
| 216 | + os.environ['SINGLESTOREDB_INIT_DB_URL'] = saved_init_url |
0 commit comments