black

Author: blinkov

tests/test_disable_discovery.py

@@ -8,10 +8,11 @@
 TEST_ERROR = "Test error"
 TEST_QUERY = "SELECT 1 + 2 AS sum"
 
+
 @pytest.fixture
 def mock_connection():
     """Mock a YDB connection to avoid actual connections."""
-    with unittest.mock.patch('ydb.connection.Connection.ready_factory') as mock_factory:
+    with unittest.mock.patch("ydb.connection.Connection.ready_factory") as mock_factory:
         # Setup the mock to return a connection-like object
         mock_connection = unittest.mock.MagicMock()
         # Use the endpoint fixture value via the function parameter
@@ -24,12 +25,12 @@ def mock_connection():
 @pytest.fixture
 def mock_aio_connection():
     """Mock a YDB async connection to avoid actual connections."""
-    with unittest.mock.patch('ydb.aio.connection.Connection.__init__') as mock_init:
+    with unittest.mock.patch("ydb.aio.connection.Connection.__init__") as mock_init:
         # Setup the mock to return None (as __init__ does)
         mock_init.return_value = None
 
         # Mock connection_ready method
-        with unittest.mock.patch('ydb.aio.connection.Connection.connection_ready') as mock_ready:
+        with unittest.mock.patch("ydb.aio.connection.Connection.connection_ready") as mock_ready:
             # Create event loop if there isn't one currently
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)
@@ -42,27 +43,29 @@ def mock_aio_connection():
 
 def create_mock_discovery_resolver(path):
     """Create a mock discovery resolver that raises exception if called."""
+
     def _mock_fixture():
         with unittest.mock.patch(path) as mock_resolve:
             # Configure mock to throw an exception if called
             mock_resolve.side_effect = Exception("Discovery should not be executed when discovery is disabled")
             yield mock_resolve
+
     return _mock_fixture
 
 
 # Mock discovery resolvers to verify no discovery requests are made
-mock_discovery_resolver = pytest.fixture(create_mock_discovery_resolver('ydb.resolver.DiscoveryEndpointsResolver.context_resolve'))
-mock_aio_discovery_resolver = pytest.fixture(create_mock_discovery_resolver('ydb.aio.resolver.DiscoveryEndpointsResolver.resolve'))
+mock_discovery_resolver = pytest.fixture(
+    create_mock_discovery_resolver("ydb.resolver.DiscoveryEndpointsResolver.context_resolve")
+)
+mock_aio_discovery_resolver = pytest.fixture(
+    create_mock_discovery_resolver("ydb.aio.resolver.DiscoveryEndpointsResolver.resolve")
+)
 
 
 # Basic unit tests for DriverConfig
 def test_driver_config_has_disable_discovery_option(endpoint, database):
     """Test that DriverConfig has the disable_discovery option."""
-    config = ydb.DriverConfig(
-        endpoint=endpoint,
-        database=database,
-        disable_discovery=True
-    )
+    config = ydb.DriverConfig(endpoint=endpoint, database=database, disable_discovery=True)
     assert hasattr(config, "disable_discovery")
     assert config.disable_discovery is True
 
@@ -110,9 +113,11 @@ def create_completed_future():
 
 
 # Mock tests for synchronous driver
-def test_sync_driver_discovery_disabled_mock(driver_config_disabled_discovery, mock_connection, mock_discovery_resolver):
+def test_sync_driver_discovery_disabled_mock(
+    driver_config_disabled_discovery, mock_connection, mock_discovery_resolver
+):
     """Test that when disable_discovery=True, the discovery thread is not started and resolver is not called (mock)."""
-    with unittest.mock.patch('ydb.pool.Discovery') as mock_discovery_class:
+    with unittest.mock.patch("ydb.pool.Discovery") as mock_discovery_class:
         driver = ydb.Driver(driver_config=driver_config_disabled_discovery)
 
         try:
@@ -129,15 +134,17 @@ def test_sync_driver_discovery_disabled_mock(driver_config_disabled_discovery, m
                 pass  # Expected exception, we just want to ensure no discovery occurs
 
             # Verify the mock wasn't called
-            assert not mock_discovery_resolver.called, "Discovery resolver should not be called when discovery is disabled"
+            assert (
+                not mock_discovery_resolver.called
+            ), "Discovery resolver should not be called when discovery is disabled"
         finally:
             # Clean up
             driver.stop()
 
 
 def test_sync_driver_discovery_enabled_mock(driver_config_enabled_discovery, mock_connection):
     """Test that when disable_discovery=False, the discovery thread is started (mock)."""
-    with unittest.mock.patch('ydb.pool.Discovery') as mock_discovery_class:
+    with unittest.mock.patch("ydb.pool.Discovery") as mock_discovery_class:
         mock_discovery_instance = unittest.mock.MagicMock()
         mock_discovery_class.return_value = mock_discovery_instance
 
@@ -158,51 +165,53 @@ def setup_async_driver_mocks():
     mocks = {}
 
     # Create mock for Discovery class
-    discovery_patcher = unittest.mock.patch('ydb.aio.pool.Discovery')
-    mocks['mock_discovery_class'] = discovery_patcher.start()
+    discovery_patcher = unittest.mock.patch("ydb.aio.pool.Discovery")
+    mocks["mock_discovery_class"] = discovery_patcher.start()
 
     # Mock the event loop
-    loop_patcher = unittest.mock.patch('asyncio.get_event_loop')
+    loop_patcher = unittest.mock.patch("asyncio.get_event_loop")
     mock_loop = loop_patcher.start()
     mock_loop_instance = unittest.mock.MagicMock()
     mock_loop.return_value = mock_loop_instance
     mock_loop_instance.create_task.return_value = unittest.mock.MagicMock()
-    mocks['mock_loop'] = mock_loop
+    mocks["mock_loop"] = mock_loop
 
     # Mock the connection pool stop method
-    stop_patcher = unittest.mock.patch('ydb.aio.pool.ConnectionPool.stop')
+    stop_patcher = unittest.mock.patch("ydb.aio.pool.ConnectionPool.stop")
     mock_stop = stop_patcher.start()
     mock_stop.return_value = create_completed_future()
-    mocks['mock_stop'] = mock_stop
+    mocks["mock_stop"] = mock_stop
 
     # Add cleanup for all patchers
-    mocks['patchers'] = [discovery_patcher, loop_patcher, stop_patcher]
+    mocks["patchers"] = [discovery_patcher, loop_patcher, stop_patcher]
 
     return mocks
 
 
 def teardown_async_mocks(mocks):
     """Clean up all mock patchers."""
-    for patcher in mocks['patchers']:
+    for patcher in mocks["patchers"]:
         patcher.stop()
 
 
 # Mock tests for asynchronous driver
 @pytest.mark.asyncio
-async def test_aio_driver_discovery_disabled_mock(driver_config_disabled_discovery, mock_aio_connection, mock_aio_discovery_resolver):
+async def test_aio_driver_discovery_disabled_mock(
+    driver_config_disabled_discovery, mock_aio_connection, mock_aio_discovery_resolver
+):
     """Test that when disable_discovery=True, the discovery is not created and resolver is not called (mock)."""
     mocks = setup_async_driver_mocks()
 
     try:
         # Mock the pool's call method to prevent unhandled exceptions
-        with unittest.mock.patch('ydb.aio.pool.ConnectionPool.__call__') as mock_call:
+        with unittest.mock.patch("ydb.aio.pool.ConnectionPool.__call__") as mock_call:
             mock_call.return_value = create_future_with_error()
 
             driver = ydb.aio.Driver(driver_config=driver_config_disabled_discovery)
 
             try:
                 # Check that the discovery class was not instantiated
-                mocks['mock_discovery_class'].assert_not_called()
+                mocks["mock_discovery_class"].assert_not_called()
 
                 # Check that discovery is disabled in debug details
                 assert_discovery_disabled(driver)
@@ -219,7 +228,9 @@ async def test_aio_driver_discovery_disabled_mock(driver_config_disabled_discove
                     pass  # Other exceptions are expected as we're using mocks
 
                 # Verify the mock wasn't called
-                assert not mock_aio_discovery_resolver.called, "Discovery resolver should not be called when discovery is disabled"
+                assert (
+                    not mock_aio_discovery_resolver.called
+                ), "Discovery resolver should not be called when discovery is disabled"
             finally:
                 # The stop method is already mocked, so we don't need to await it
                 pass
@@ -234,13 +245,13 @@ async def test_aio_driver_discovery_enabled_mock(driver_config_enabled_discovery
 
     try:
         mock_discovery_instance = unittest.mock.MagicMock()
-        mocks['mock_discovery_class'].return_value = mock_discovery_instance
+        mocks["mock_discovery_class"].return_value = mock_discovery_instance
 
         driver = ydb.aio.Driver(driver_config=driver_config_enabled_discovery)
 
         try:
             # Check that the discovery class was instantiated
-            mocks['mock_discovery_class'].assert_called_once()
+            mocks["mock_discovery_class"].assert_called_once()
         finally:
             # The stop method is already mocked, so we don't need to await it
             pass
@@ -265,6 +276,7 @@ def test_integration_disable_discovery(driver_config_disabled_discovery):
 
         # Try to execute a simple query to ensure it works with discovery disabled
         with ydb.SessionPool(driver) as pool:
+
             def query_callback(session):
                 result_sets = session.transaction().execute(TEST_QUERY, commit_tx=True)
                 assert len(result_sets) == 1
@@ -297,4 +309,4 @@ async def query_callback(session):
         finally:
             await session_pool.stop()
     finally:
-        await driver.stop(timeout=10)
+        await driver.stop(timeout=10)