Added getitem method and test (#29)

exasol · May 10, 2022 · 30e18e8 · 30e18e8
1 parent 008ba58
commit 30e18e8
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 3 deletions.
diff --git a/exasol_udf_mock_python/mock_context_run_wrapper.py b/exasol_udf_mock_python/mock_context_run_wrapper.py
@@ -3,11 +3,14 @@
 
 def _disallowed_function(*args, **kw):
     raise RuntimeError(
-        "F-UDF-CL-SL-PYTHON-1107: next(), reset() and emit() functions are not allowed in scalar context")
+        "F-UDF-CL-SL-PYTHON-1107: next(), reset() and emit() "
+        "functions are not allowed in scalar context")
+
 
 class MockContextRunWrapper:
 
-    def __init__(self, mock_context: MockContext, input_type: str, output_type: str):
+    def __init__(
+            self, mock_context: MockContext, input_type: str, output_type: str):
         self._output_type = output_type
         self._input_type = input_type
         self._mock_context = mock_context
@@ -24,6 +27,8 @@ def __init__(self, mock_context: MockContext, input_type: str, output_type: str)
             self.get_dataframe = self._mock_context.get_dataframe
             self.size = self._mock_context.size
 
-
     def __getattr__(self, name):
         return self._mock_context.__getattr__(name)
+
+    def __getitem__(self, item):
+        return self._mock_context._data[item]
diff --git a/tests/test_executor_context_set_emits.py b/tests/test_executor_context_set_emits.py
@@ -364,3 +364,34 @@ def run(ctx):
     exa = MockExaEnvironment(meta)
     with pytest.raises(TypeError):
         result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+
+
+def test_context_parameters():
+    def udf_wrapper():
+        def run(ctx):
+            ctx.emit(ctx[0], ctx.t1)
+            ctx.emit(ctx[1], ctx.t2)
+            ctx.emit(ctx[2], ctx.t3)
+
+    input_columns = [Column("t1", int, "INTEGER"),
+                     Column("t2", int, "INTEGER"),
+                     Column("t3", int, "INTEGER")]
+    output_columns = [Column("o1", int, "INTEGER"),
+                      Column("o2", int, "INTEGER")]
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=input_columns,
+        output_type="EMITS",
+        output_columns=output_columns
+    )
+    input_data = [(1, 2, 3), (4, 5, 6)]
+    exa = MockExaEnvironment(meta)
+    executor = UDFMockExecutor()
+    result = executor.run([Group(input_data)], exa)
+    for i, group in enumerate(result):
+        result_row = group.rows
+        assert len(result_row) == len(input_columns)
+        for j in range(len(result_row)):
+            assert len(result_row[j]) == len(output_columns)
+            assert input_data[i][j] == result_row[j][0] == result_row[j][1]