✨ add force_pagination_size and update docs (#30)

* ✨ add force_pagination_size and update docs

* 📝 fix docs bug

* 🔖 Bump version number

---------

Co-authored-by: itsthejoker <itsthejoker@users.noreply.github.com>

asbestos/asbestos.py

@@ -40,12 +40,14 @@ def __init__(
         response: dict[Any],
         ephemeral: bool = False,
         data: Optional[tuple[Any]] = None,
+        force_pagination_size: Optional[int] = None,
     ) -> None:
         self.query: str = query
         self.data: Optional[tuple[Any]] = data
         self.response: dict[Any] | list[dict[Any]] = response
         self.ephemeral: bool = ephemeral
         self.sfqid: int = 0
+        self.force_pagination_size = force_pagination_size
         self.set_sfqid()
 
     def __str__(self) -> str:
@@ -143,7 +145,13 @@ def check_for_duplicates(self, query: str, data: tuple, ephemeral: bool) -> None
                     )
                 raise AsbestosDuplicateQuery("You've already registered this query!")
 
-    def register(self, query: str, response: dict, data: tuple = None) -> None:
+    def register(
+        self,
+        query: str,
+        response: dict,
+        data: tuple = None,
+        force_pagination_size: int = None,
+    ) -> None:
         """
         Asbestos will watch for any call that includes the query passed in
         and return the response you save. If data is provided with the query,
@@ -183,15 +191,42 @@ def register(self, query: str, response: dict, data: tuple = None) -> None:
         >>> cursor.execute("A", ("E",))  # extra data doesn't match!
         >>> cursor.fetchone()
         {}  # default response
+        ```
+
+        If you're dealing with a `fetchmany()` call and need to control
+        how many results are returned regardless of what the code requests,
+        pass the return count into `force_pagination_size`. This will override
+        both `cursor.arraysize` and the number passed into `fetchmany`.
+
+        ```python
+        >>> config.register(
+        ...     query="A",
+        ...     response=[{'a': 1}, {'b': 2}],
+        ...     force_pagination_size=1
+        ... )
+        >>> cursor.execute("A")
+        >>> cursor.fetchmany(50)
+        {'a': 1}  # because we overrode the value when registering
+
         ```
         """
         self.check_for_duplicates(query=query, data=data, ephemeral=False)
         self.query_map.append(
-            AsbestosResponse(query=query, data=data, response=response, ephemeral=False)
+            AsbestosResponse(
+                query=query,
+                data=data,
+                response=response,
+                ephemeral=False,
+                force_pagination_size=force_pagination_size,
+            )
         )
 
     def register_ephemeral(
-        self, query: str, response: dict, data: tuple = None
+        self,
+        query: str,
+        response: dict,
+        data: tuple = None,
+        force_pagination_size: int = None,
     ) -> None:
         """
         Works the same way as `AsbestosConfig.register()` with the only difference being
@@ -200,7 +235,13 @@ def register_ephemeral(
         """
         self.check_for_duplicates(query=query, data=data, ephemeral=True)
         self.query_map.append(
-            AsbestosResponse(query=query, data=data, response=response, ephemeral=True)
+            AsbestosResponse(
+                query=query,
+                data=data,
+                response=response,
+                ephemeral=True,
+                force_pagination_size=force_pagination_size,
+            )
         )
 
     def clear_queries(self) -> None:
@@ -299,7 +340,7 @@ def fetchone(self) -> dict:
         config.register(query="A", response=[{'a':1}, {'b': 2}])
         with asbestos_cursor() as cursor:
             cursor.execute("A")
-            resp = cursor.fetchall()
+            resp = cursor.fetchone()
 
         # resp = {'a':1}
         ```
@@ -342,16 +383,22 @@ def fetchmany(self, size: int = None) -> list[dict]:
         cursor.arraysize = 1  # return one response at a time
         with asbestos_cursor() as cursor:
             cursor.execute("A")
-            cursor.fetchmany(1)  # [{'a':1}]
-            cursor.fetchmany(1)  # [{'b': 2}]
-            cursor.fetchmany(1)  # []
+            cursor.fetchmany()  # [{'a':1}]
+            cursor.fetchmany()  # [{'b': 2}]
+            cursor.fetchmany()  # []
+
+        Note: a value passed directly into `fetchmany` overrides the
+        `cursor.arraysize`, and both of them are overridden by registering
+        the query with `force_pagination_size`.
         ```
         """
         size = size if size else self.arraysize
         if self.last_paginated_query != self.config.last_run_query.sfqid:
             # this is the first time we're paginating here
             self.last_page_start = 0
-            self.last_paginated_query = self.config.last_run_query.sfqid
+            self.last_paginated_query: int = self.config.last_run_query.sfqid
+        if query_forced_page_size := self.config.last_run_query.force_pagination_size:
+            size = query_forced_page_size
 
         response = self.config.last_run_query.response[
             self.last_page_start : size + self.last_page_start

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "asbestos-snow"
-version = "1.0.4"
+version = "1.1.0"
 description = "An easy way to mock Snowflake connections in Python!"
 authors = ["Joe Kaufeld <jkaufeld@spoton.com>", "SpotOn <opensource@spoton.com>"]
 readme = "README.md"

tests/test_asbestos.py

@@ -7,6 +7,30 @@
 QUERY = "query"
 INVALID_QUERY = "asdf"
 RESPONSE = [{"response": "hi"}, {"hello there": "general kenobi"}]
+SHORT_BATCH_RESPONSE = [
+    {"a": 1},
+    {"b": 2},
+    {"c": 3},
+    {"d": 4},
+    {"e": 5},
+]
+LARGE_BATCH_RESPONSE = [
+    {"a": 1},
+    {"b": 2},
+    {"c": 3},
+    {"d": 4},
+    {"e": 5},
+    {"f": 6},
+    {"g": 7},
+    {"h": 8},
+    {"i": 9},
+    {"j": 10},
+    {"k": 11},
+    {"l": 12},
+    {"m": 13},
+    {"n": 14},
+    {"o": 15},
+]
 
 
 @pytest.fixture(autouse=True)
@@ -174,13 +198,7 @@ def test_async_query_registered_response() -> None:
 def test_fetchmany() -> None:
     config.register(
         query=QUERY,
-        response=[
-            {"a": 1},
-            {"b": 2},
-            {"c": 3},
-            {"d": 4},
-            {"e": 5},
-        ],
+        response=SHORT_BATCH_RESPONSE,
     )
     with asbestos_cursor() as cursor:
         cursor.arraysize = 2
@@ -190,11 +208,58 @@ def test_fetchmany() -> None:
         assert cursor.fetchmany() == [{"e": 5}]
 
 
+def test_fetchmany_with_value_in_call() -> None:
+    config.register(
+        query=QUERY,
+        response=SHORT_BATCH_RESPONSE,
+    )
+    with asbestos_cursor() as cursor:
+        cursor.execute(QUERY)
+        assert cursor.fetchmany(2) == [{"a": 1}, {"b": 2}]
+        assert cursor.fetchmany(2) == [{"c": 3}, {"d": 4}]
+        assert cursor.fetchmany(2) == [{"e": 5}]
+
+
 def test_fetchmany_default_page() -> None:
-    config.register(query=QUERY, response=RESPONSE)
+    config.register(query=QUERY, response=LARGE_BATCH_RESPONSE)
+    with asbestos_cursor() as cursor:
+        cursor.execute(QUERY)
+        assert cursor.fetchmany() == LARGE_BATCH_RESPONSE[:10]
+
+
+def test_fetchmany_local_size_overrides_arraysize() -> None:
+    config.register(
+        query=QUERY,
+        response=SHORT_BATCH_RESPONSE,
+    )
+    with asbestos_cursor() as cursor:
+        cursor.execute(QUERY)
+        cursor.arraysize = 10
+        assert cursor.fetchmany(2) == [{"a": 1}, {"b": 2}]
+        assert cursor.fetchmany(2) == [{"c": 3}, {"d": 4}]
+        assert cursor.fetchmany(2) == [{"e": 5}]
+
+
+def test_fetchmany_force_pagination_size() -> None:
+    config.register(query=QUERY, response=SHORT_BATCH_RESPONSE, force_pagination_size=2)
+    with asbestos_cursor() as cursor:
+        cursor.execute(QUERY)
+        cursor.arraysize = 10
+        assert cursor.fetchmany(5) == [{"a": 1}, {"b": 2}]
+        assert cursor.fetchmany(5) == [{"c": 3}, {"d": 4}]
+        assert cursor.fetchmany(5) == [{"e": 5}]
+
+
+def test_fetchmany_force_pagination_size_ephemeral() -> None:
+    config.register_ephemeral(
+        query=QUERY, response=SHORT_BATCH_RESPONSE, force_pagination_size=2
+    )
     with asbestos_cursor() as cursor:
         cursor.execute(QUERY)
-        assert cursor.fetchmany() == RESPONSE
+        cursor.arraysize = 10
+        assert cursor.fetchmany(5) == [{"a": 1}, {"b": 2}]
+        assert cursor.fetchmany(5) == [{"c": 3}, {"d": 4}]
+        assert cursor.fetchmany(5) == [{"e": 5}]
 
 
 def test_documentation_example() -> None: