additional unit tests for the new contextualization mechanism

Author: ds-jakub-cierocki

tests/unit/views/test_methods_base.py

@@ -1,13 +1,28 @@
 # pylint: disable=missing-docstring, missing-return-doc, missing-param-doc, disallowed-name
 
-
-from typing import List, Literal, Tuple
+import asyncio
+from dataclasses import dataclass
+from typing import List, Literal, Tuple, Union, Optional
 
 from dbally.collection.results import ViewExecutionResult
 from dbally.iql import IQLQuery
 from dbally.views.decorators import view_filter
-from dbally.views.exposed_functions import MethodParamWithTyping
+from dbally.views.exposed_functions import MethodParamWithTyping, ExposedFunction
 from dbally.views.methods_base import MethodsBaseView
+from dbally.context import BaseCallerContext
+from dbally.iql_generator.iql_generator import IQLGenerator
+from dbally.audit.event_tracker import EventTracker
+from dbally.prompt.elements import FewShotExample
+from dbally.llms.clients.base import LLMOptions
+from dbally.llms.base import LLM
+
+
+@dataclass
+class TestCallerContext(BaseCallerContext):
+    """
+    Mock class for testing context.
+    """
+    current_year: Literal['2023', '2024']
 
 
 class MockMethodsBase(MethodsBaseView):
@@ -22,7 +37,7 @@ def method_foo(self, idx: int) -> None:
         """
 
     @view_filter()
-    def method_bar(self, cities: List[str], year: Literal["2023", "2024"], pairs: List[Tuple[str, int]]) -> str:
+    def method_bar(self, cities: List[str], year: Union[Literal["2023", "2024"], TestCallerContext], pairs: List[Tuple[str, int]]) -> str:
         return f"hello {cities} in {year} of {pairs}"
 
     async def apply_filters(self, filters: IQLQuery) -> None:
@@ -47,9 +62,24 @@ def test_list_filters() -> None:
     assert method_bar.description == ""
     assert method_bar.parameters == [
         MethodParamWithTyping("cities", List[str]),
-        MethodParamWithTyping("year", Literal["2023", "2024"]),
+        MethodParamWithTyping("year", Union[Literal["2023", "2024"], TestCallerContext]),
         MethodParamWithTyping("pairs", List[Tuple[str, int]]),
     ]
     assert (
-        str(method_bar) == "method_bar(cities: List[str], year: Literal['2023', '2024'], pairs: List[Tuple[str, int]])"
+        str(method_bar) == "method_bar(cities: List[str], year: Literal['2023', '2024'] | AskerContext, pairs: List[Tuple[str, int]])"
     )
+
+
+async def test_contextualization() -> None:
+    mock_view = MockMethodsBase()
+    filters = mock_view.list_filters()
+    test_context = TestCallerContext("2024")
+    mock_view.contextualize_filters(filters, [test_context])
+
+    method_foo = [f for f in filters if f.name == "method_foo"][0]
+    assert method_foo.context_class is None
+    assert method_foo.context is None
+
+    method_bar = [f for f in filters if f.name == "method_bar"][0]
+    assert method_bar.context_class is TestCallerContext
+    assert method_bar.context is test_context