test: reduce query size for v1 long test (#339)

src/firebolt/service/V1/base.py

@@ -1,11 +1,13 @@
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from firebolt.client import ClientV1 as Client
-from firebolt.service.manager import ResourceManager
+
+if TYPE_CHECKING:
+    from firebolt.service.manager import ResourceManager
 
 
 class BaseService:
-    def __init__(self, resource_manager: ResourceManager):
+    def __init__(self, resource_manager: "ResourceManager"):
         self.resource_manager = resource_manager
 
     @property

src/firebolt/service/V1/region.py

@@ -1,14 +1,16 @@
-from typing import Dict, List
+from typing import TYPE_CHECKING, Dict, List
 
 from firebolt.model.V1.region import Region, RegionKey
-from firebolt.service.manager import ResourceManager
 from firebolt.service.V1.base import BaseService
 from firebolt.utils.urls import REGIONS_URL
 from firebolt.utils.util import cached_property
 
+if TYPE_CHECKING:
+    from firebolt.service.manager import ResourceManager
+
 
 class RegionService(BaseService):
-    def __init__(self, resource_manager: ResourceManager):
+    def __init__(self, resource_manager: "ResourceManager"):
         """
         Service to manage AWS regions (us-east-1, etc)
 

src/firebolt/service/manager.py

@@ -14,7 +14,14 @@
 )
 from firebolt.common import Settings
 from firebolt.db import connect
+from firebolt.service.V1.binding import BindingService
+from firebolt.service.V1.database import DatabaseService as DatabaseServiceV1
+from firebolt.service.V1.engine import EngineService as EngineServiceV1
 from firebolt.service.V1.provider import get_provider_id
+from firebolt.service.V1.region import RegionService
+from firebolt.service.V2.database import DatabaseService as DatabaseServiceV2
+from firebolt.service.V2.engine import EngineService as EngineServiceV2
+from firebolt.service.V2.instance_type import InstanceTypeService
 from firebolt.utils.util import fix_url_schema
 
 DEFAULT_TIMEOUT_SECONDS: int = 60 * 2
@@ -127,34 +134,25 @@ def __init__(
 
     def _init_services_v2(self) -> None:
         # avoid circular import
-        from firebolt.service.V2.database import DatabaseService
-        from firebolt.service.V2.engine import EngineService
-        from firebolt.service.V2.instance_type import InstanceTypeService
 
         # Cloud Platform Resources (AWS)
         self.instance_types = InstanceTypeService(resource_manager=self)
 
         # Firebolt Resources
-        self.databases = DatabaseService(resource_manager=self)
-        self.engines = EngineService(resource_manager=self)
+        self.databases = DatabaseServiceV2(resource_manager=self)
+        self.engines = EngineServiceV2(resource_manager=self)
 
         # Not applicable to V2
         self.provider_id = None
 
     def _init_services_v1(self) -> None:
-        # avoid circular import
-        from firebolt.service.V1.binding import BindingService
-        from firebolt.service.V1.database import DatabaseService
-        from firebolt.service.V1.engine import EngineService
-        from firebolt.service.V1.region import RegionService
-
         # Cloud Platform Resources (AWS)
         self.regions = RegionService(resource_manager=self)  # type: ignore
 
         # Firebolt Resources
         self.bindings = BindingService(resource_manager=self)  # type: ignore
-        self.engines = EngineService(resource_manager=self)  # type: ignore
-        self.databases = DatabaseService(resource_manager=self)  # type: ignore
+        self.engines = EngineServiceV1(resource_manager=self)  # type: ignore
+        self.databases = DatabaseServiceV1(resource_manager=self)  # type: ignore
 
         self.provider_id = get_provider_id(client=self._client)
 

tests/integration/conftest.py

@@ -1,5 +1,7 @@
 from logging import getLogger
 from os import environ
+from time import time
+from typing import Optional
 
 from pytest import fixture, mark
 
@@ -142,3 +144,20 @@ def engine_url() -> str:
 @fixture(scope="session")
 def stopped_engine_url() -> str:
     return must_env(STOPPED_ENGINE_URL_ENV)
+
+
+@fixture(scope="function")
+def minimal_time():
+    limit: Optional[float] = None
+
+    def setter(value):
+        nonlocal limit
+        limit = value
+
+    start = time()
+    yield setter
+    end = time()
+    if limit is not None:
+        assert (
+            end - start >= limit
+        ), f"Test took {end - start} seconds, less than {limit} seconds"

tests/integration/dbapi/async/V1/test_queries_async.py

@@ -1,6 +1,6 @@
 from datetime import date, datetime
 from decimal import Decimal
-from typing import Any, List
+from typing import Any, Callable, List
 
 from pytest import fixture, mark, raises
 
@@ -170,14 +170,19 @@ async def test_select(
 
 
 @mark.slow
-@mark.timeout(timeout=600)
+@mark.timeout(timeout=1000)
 async def test_long_query(
     connection: Connection,
+    minimal_time: Callable[[float], None],
 ) -> None:
     """AWS ALB TCP timeout set to 350; make sure we handle the keepalive correctly."""
+
+    # Fail test if it takes less than 350 seconds
+    minimal_time(350)
+
     with connection.cursor() as c:
         await c.execute(
-            "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)",  # approx 6m runtime
+            "SELECT checksum(*) FROM GENERATE_SERIES(1, 250000000000)",  # approx 6m runtime
         )
         data = await c.fetchall()
         assert len(data) == 1, "Invalid data size returned by fetchall"

tests/integration/dbapi/async/V2/test_queries_async.py

@@ -1,7 +1,7 @@
 from datetime import date, datetime
 from decimal import Decimal
 from os import environ
-from typing import List
+from typing import Callable, List
 
 from pytest import fixture, mark, raises
 
@@ -104,8 +104,12 @@ async def test_select(
 @mark.timeout(timeout=550)
 async def test_long_query(
     connection: Connection,
+    minimal_time: Callable[[float], None],
 ) -> None:
     """AWS ALB TCP timeout set to 350; make sure we handle the keepalive correctly."""
+
+    minimal_time(350)
+
     with connection.cursor() as c:
         await c.execute(
             "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)",  # approx 6m runtime

tests/integration/dbapi/sync/V1/test_queries.py

@@ -1,7 +1,7 @@
 from datetime import date, datetime
 from decimal import Decimal
 from threading import Thread
-from typing import Any, List
+from typing import Any, Callable, List
 
 from pytest import fixture, mark, raises
 
@@ -121,14 +121,18 @@ def test_select(
 
 
 @mark.slow
-@mark.timeout(timeout=600)
+@mark.timeout(timeout=1000)
 def test_long_query(
-    connection: Connection,
+    connection: Connection, minimal_time: Callable[[float], None]
 ) -> None:
     """AWS ALB TCP timeout set to 350, make sure we handle the keepalive correctly."""
+
+    # Fail test if it takes less than 350 seconds
+    minimal_time(350)
+
     with connection.cursor() as c:
         c.execute(
-            "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)",  # approx 6m runtime
+            "SELECT checksum(*) FROM GENERATE_SERIES(1, 250000000000)",  # approx 6m runtime
         )
         data = c.fetchall()
         assert len(data) == 1, "Invalid data size returned by fetchall"

tests/integration/dbapi/sync/V2/test_queries.py

@@ -2,7 +2,7 @@
 from decimal import Decimal
 from os import environ
 from threading import Thread
-from typing import Any, List
+from typing import Any, Callable, List
 
 from pytest import fixture, mark, raises
 
@@ -110,8 +110,12 @@ def test_select(
 @mark.timeout(timeout=550)
 def test_long_query(
     connection: Connection,
+    minimal_time: Callable[[float], None],
 ) -> None:
-    """AWS ALB TCP timeout set to 350, make sure we handle the keepalive correctly."""
+    """AWS ALB TCP timeout set to 350; make sure we handle the keepalive correctly."""
+
+    minimal_time(350)
+
     with connection.cursor() as c:
         c.execute(
             "SELECT checksum(*) FROM GENERATE_SERIES(1, 200000000000)",  # approx 6m runtime