Skip to content

Commit 30e18e8

Browse files
Added getitem method and test (#29)
1 parent 008ba58 commit 30e18e8

File tree

2 files changed

+39
-3
lines changed

2 files changed

+39
-3
lines changed

exasol_udf_mock_python/mock_context_run_wrapper.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,14 @@
33

44
def _disallowed_function(*args, **kw):
55
raise RuntimeError(
6-
"F-UDF-CL-SL-PYTHON-1107: next(), reset() and emit() functions are not allowed in scalar context")
6+
"F-UDF-CL-SL-PYTHON-1107: next(), reset() and emit() "
7+
"functions are not allowed in scalar context")
8+
79

810
class MockContextRunWrapper:
911

10-
def __init__(self, mock_context: MockContext, input_type: str, output_type: str):
12+
def __init__(
13+
self, mock_context: MockContext, input_type: str, output_type: str):
1114
self._output_type = output_type
1215
self._input_type = input_type
1316
self._mock_context = mock_context
@@ -24,6 +27,8 @@ def __init__(self, mock_context: MockContext, input_type: str, output_type: str)
2427
self.get_dataframe = self._mock_context.get_dataframe
2528
self.size = self._mock_context.size
2629

27-
2830
def __getattr__(self, name):
2931
return self._mock_context.__getattr__(name)
32+
33+
def __getitem__(self, item):
34+
return self._mock_context._data[item]

tests/test_executor_context_set_emits.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -364,3 +364,34 @@ def run(ctx):
364364
exa = MockExaEnvironment(meta)
365365
with pytest.raises(TypeError):
366366
result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
367+
368+
369+
def test_context_parameters():
370+
def udf_wrapper():
371+
def run(ctx):
372+
ctx.emit(ctx[0], ctx.t1)
373+
ctx.emit(ctx[1], ctx.t2)
374+
ctx.emit(ctx[2], ctx.t3)
375+
376+
input_columns = [Column("t1", int, "INTEGER"),
377+
Column("t2", int, "INTEGER"),
378+
Column("t3", int, "INTEGER")]
379+
output_columns = [Column("o1", int, "INTEGER"),
380+
Column("o2", int, "INTEGER")]
381+
meta = MockMetaData(
382+
script_code_wrapper_function=udf_wrapper,
383+
input_type="SET",
384+
input_columns=input_columns,
385+
output_type="EMITS",
386+
output_columns=output_columns
387+
)
388+
input_data = [(1, 2, 3), (4, 5, 6)]
389+
exa = MockExaEnvironment(meta)
390+
executor = UDFMockExecutor()
391+
result = executor.run([Group(input_data)], exa)
392+
for i, group in enumerate(result):
393+
result_row = group.rows
394+
assert len(result_row) == len(input_columns)
395+
for j in range(len(result_row)):
396+
assert len(result_row[j]) == len(output_columns)
397+
assert input_data[i][j] == result_row[j][0] == result_row[j][1]

0 commit comments

Comments
 (0)