Add checks for parameter of get_dataframe (#20)

Author: tkilias

exasol_udf_mock_python/mock_context.py

@@ -14,10 +14,10 @@ def __init__(self, input_groups: Iterator[Group], metadata: MockMetaData):
         self._input_groups = input_groups
         self._output_groups = []
         self._input_group = None  # type: Group
-        self._output_group_list = None # type: List
+        self._output_group_list = None  # type: List
         self._output_group = None  # type: Group
         self._iter = None  # type: Iterator[Tuple]
-        self._len = None # type: int
+        self._len = None  # type: int
         self._metadata = metadata
         self._name_position_map = \
             {column.name: position
@@ -51,7 +51,14 @@ def _next_group(self):
         self.next()
         return True
 
+    def _is_positive_integer(self, value):
+        return value is not None and isinstance(value, int) and value > 0
+
     def get_dataframe(self, num_rows='all', start_col=0):
+        if not (num_rows == 'all' or self._is_positive_integer(num_rows)):
+            raise RuntimeError("get_dataframe() parameter 'num_rows' must be 'all' or an integer > 0")
+        if not (self._is_positive_integer(start_col) or start_col == 0):
+            raise RuntimeError("get_dataframe() parameter 'start_col' must be an integer >= 0")
         if self._data is None:
             return None
         columns_ = [column.name for column in self._metadata.input_columns]
@@ -66,15 +73,15 @@ def get_dataframe(self, num_rows='all', start_col=0):
                 df = df.append(df_current)
             if not self.next():
                 break
-            i+=1
+            i += 1
         if df is not None:
             df = df.reset_index(drop=True)
         return df
 
     def __getattr__(self, name):
         return self._data[self._name_position_map[name]]
 
-    def next(self, reset:bool = False):
+    def next(self, reset: bool = False):
         if reset:
             self.reset()
         else:

tests/test_executor_context_set_emits.py

@@ -28,9 +28,9 @@ def run(ctx):
     result = executor.run([Group([(1,), (5,), (6,)])], exa)
     assert result == [Group([(1,), (5,), (6,)])]
 
+
 def test_emit_single_column_none():
     def udf_wrapper():
-
         def run(ctx):
             ctx.emit(None)
 
@@ -46,11 +46,11 @@ def run(ctx):
     result = executor.run([Group([(1,), (5,), (6,)])], exa)
     assert result == [Group([(None,)])]
 
+
 def test_emit_multi_column_none():
     def udf_wrapper():
-
         def run(ctx):
-            ctx.emit(None,None)
+            ctx.emit(None, None)
 
     executor = UDFMockExecutor()
     meta = MockMetaData(
@@ -63,7 +63,8 @@ def run(ctx):
     )
     exa = MockExaEnvironment(meta)
     result = executor.run([Group([(1,), (5,), (6,)])], exa)
-    assert result == [Group([(None,None)])]
+    assert result == [Group([(None, None)])]
+
 
 def test_next_emit_reset():
     def udf_wrapper():
@@ -75,7 +76,7 @@ def run(ctx):
                     break
             ctx.reset()
             while True:
-                ctx.emit(ctx.t+1)
+                ctx.emit(ctx.t + 1)
                 if not ctx.next():
                     break
 
@@ -91,6 +92,7 @@ def run(ctx):
     result = executor.run([Group([(1,), (5,), (6,)])], exa)
     assert result == [Group([(1,), (5,), (6,), (2,), (6,), (7,)])]
 
+
 def test_next_reset_combined():
     def udf_wrapper():
 
@@ -101,7 +103,7 @@ def run(ctx):
                     break
             ctx.next(reset=True)
             for i in range(2):
-                ctx.emit(ctx.t+1)
+                ctx.emit(ctx.t + 1)
                 if not ctx.next():
                     break
 
@@ -115,7 +117,7 @@ def run(ctx):
     )
     exa = MockExaEnvironment(meta)
     result = executor.run([Group([(1,), (5,), (6,)])], exa)
-    assert result == [Group([(1,), (5,),(2,), (6,)])]
+    assert result == [Group([(1,), (5,), (2,), (6,)])]
 
 
 def test_get_dataframe_all():
@@ -185,9 +187,168 @@ def run(ctx):
     result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
     assert result == [Group([(1,), (2,), (4,), (5,)])]
 
-def test_emit_tuple_exception():
+
+def test_get_dataframe_num_rows_1():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=1)
+            ctx.emit(df)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+    assert result == [Group([(1,), ])]
+
+
+def test_get_dataframe_num_rows_0():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=0)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    with pytest.raises(RuntimeError) as excinfo:
+        result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+
+
+def test_get_dataframe_num_rows_float():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=1.5)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    with pytest.raises(RuntimeError) as excinfo:
+        result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+
+def test_get_dataframe_num_rows_None():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=None)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    with pytest.raises(RuntimeError) as excinfo:
+        result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+
+
+def test_get_dataframe_num_rows_negative():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=-1)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    with pytest.raises(RuntimeError) as excinfo:
+        result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+
+def test_get_dataframe_start_col_None():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=10, start_col=None)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    with pytest.raises(RuntimeError) as excinfo:
+        result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+
+def test_get_dataframe_start_col_negative():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=10, start_col=-1)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    with pytest.raises(RuntimeError) as excinfo:
+        result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+
+def test_get_dataframe_start_col_0():
     def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=1, start_col=0)
+            ctx.emit(df)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+    assert result == [Group([(1,), ])]
 
+def test_get_dataframe_start_col_positive():
+    def udf_wrapper():
+        def run(ctx):
+            df = ctx.get_dataframe(num_rows=1, start_col=1)
+            ctx.emit(df)
+
+    executor = UDFMockExecutor()
+    meta = MockMetaData(
+        script_code_wrapper_function=udf_wrapper,
+        input_type="SET",
+        input_columns=[Column("t", int, "INTEGER")],
+        output_type="EMITS",
+        output_columns=[Column("t", int, "INTEGER")]
+    )
+    exa = MockExaEnvironment(meta)
+    result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
+    assert result == [Group([(1,), ])]
+
+def test_emit_tuple_exception():
+    def udf_wrapper():
         def run(ctx):
             while True:
                 ctx.emit((1,))