diff --git a/exasol_udf_mock_python/mock_context_run_wrapper.py b/exasol_udf_mock_python/mock_context_run_wrapper.py index fe5536a..66178cd 100644 --- 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 index 9ccac94..fe725f1 100644 --- 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]