contextualization mechanism - dedicated unit tests

Author: ds-jakub-cierocki

tests/unit/iql/test_iql_parser.py

@@ -1,28 +1,44 @@
 import re
-from typing import List
+from typing import List, Union
 
 import pytest
 
+from dbally.context import BaseCallerContext
 from dbally.iql import IQLArgumentParsingError, IQLQuery, IQLUnsupportedSyntaxError, syntax
 from dbally.iql._exceptions import IQLArgumentValidationError, IQLFunctionNotExists
 from dbally.iql._processor import IQLProcessor
 from dbally.views.exposed_functions import ExposedFunction, MethodParamWithTyping
 
 
+class TestCustomContext(BaseCallerContext):
+    city: str
+
+
+class AnotherTestCustomContext(BaseCallerContext):
+    some_field: str
+
+
 async def test_iql_parser():
+    custom_context = TestCustomContext(city="cracow")
+    custom_context2 = AnotherTestCustomContext(some_field="aaa")
+
     parsed = await IQLQuery.parse(
-        "not (filter_by_name(['John', 'Anne']) and filter_by_city('cracow') and filter_by_company('deepsense.ai'))",
+        "not (filter_by_name(['John', 'Anne']) and filter_by_city(BaseCallerContext()) and filter_by_company('deepsense.ai'))",
         allowed_functions=[
             ExposedFunction(
                 name="filter_by_name", description="", parameters=[MethodParamWithTyping(name="name", type=List[str])]
             ),
             ExposedFunction(
-                name="filter_by_city", description="", parameters=[MethodParamWithTyping(name="city", type=str)]
+                name="filter_by_city",
+                description="",
+                parameters=[MethodParamWithTyping(name="city", type=Union[str, TestCustomContext])],
+                context_class=TestCustomContext,
             ),
             ExposedFunction(
                 name="filter_by_company", description="", parameters=[MethodParamWithTyping(name="company", type=str)]
             ),
         ],
+        contexts=[custom_context, custom_context2],
     )
 
     not_op = parsed.root
@@ -37,7 +53,7 @@ async def test_iql_parser():
     assert name_filter.arguments[0] == ["John", "Anne"]
 
     assert isinstance(city_filter, syntax.FunctionCall)
-    assert city_filter.arguments[0] == "cracow"
+    assert city_filter.arguments[0] is custom_context
 
     assert isinstance(company_filter, syntax.FunctionCall)
     assert company_filter.arguments[0] == "deepsense.ai"

tests/unit/views/test_sqlalchemy_base.py

@@ -1,14 +1,20 @@
 # pylint: disable=missing-docstring, missing-return-doc, missing-param-doc, disallowed-name
 
 import re
+from typing import Union
 
 import sqlalchemy
 
+from dbally.context import BaseCallerContext
 from dbally.iql import IQLQuery
 from dbally.views.decorators import view_filter
 from dbally.views.sqlalchemy_base import SqlAlchemyBaseView
 
 
+class SomeTestContext(BaseCallerContext):
+    age: int
+
+
 class MockSqlAlchemyView(SqlAlchemyBaseView):
     """
     Mock class for testing the SqlAlchemyBaseView
@@ -22,12 +28,20 @@ def method_foo(self, idx: int) -> sqlalchemy.ColumnElement:
         """
         Some documentation string
         """
+
         return sqlalchemy.literal(idx)
 
     @view_filter()
     async def method_bar(self, city: str, year: int) -> sqlalchemy.ColumnElement:
         return sqlalchemy.literal(f"hello {city} in {year}")
 
+    @view_filter()
+    async def method_baz(self, age: Union[int, SomeTestContext]) -> sqlalchemy.ColumnElement:
+        if isinstance(age, SomeTestContext):
+            return sqlalchemy.literal(age.age)
+
+        return sqlalchemy.literal(age)
+
 
 def normalize_whitespace(s: str) -> str:
     """
@@ -44,9 +58,10 @@ async def test_filter_sql_generation() -> None:
     mock_connection = sqlalchemy.create_mock_engine("postgresql://", executor=None)
     mock_view = MockSqlAlchemyView(mock_connection.engine)
     query = await IQLQuery.parse(
-        'method_foo(1) and method_bar("London", 2020)',
+        'method_foo(1) and method_bar("London", 2020) and method_baz(BaseCallerContext())',
         allowed_functions=mock_view.list_filters(),
+        contexts=[SomeTestContext(age=69)],
     )
     await mock_view.apply_filters(query)
     sql = normalize_whitespace(mock_view.execute(dry_run=True).context["sql"])
-    assert sql == "SELECT 'test' AS foo WHERE 1 AND 'hello London in 2020'"
+    assert sql == "SELECT 'test' AS foo WHERE 1 AND 'hello London in 2020' AND 69"