Add unit test for a wrong type of output format in table prediction

Unstructured-IO · Apr 16, 2024 · b14d6e8 · b14d6e8
1 parent b23fc50
commit b14d6e8
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/test_unstructured_inference/models/test_tables.py b/test_unstructured_inference/models/test_tables.py
@@ -932,6 +932,25 @@ def test_table_prediction_output_format(
         assert expectation in result
 
 
+def test_table_prediction_output_format_when_wrong_type_then_value_error(
+    table_transformer,
+    example_image,
+    mocker,
+    example_table_cells,
+    mocked_ocr_tokens,
+):
+    mocker.patch.object(tables, "recognize", return_value=example_table_cells)
+    mocker.patch.object(
+        tables.UnstructuredTableTransformerModel,
+        "get_structure",
+        return_value=None,
+    )
+    with pytest.raises(ValueError):
+        table_transformer.run_prediction(
+            example_image, result_format="Wrong format", ocr_tokens=mocked_ocr_tokens
+        )
+
+
 def test_table_prediction_runs_with_empty_recognize(
     table_transformer,
     example_image,