feat: v2 client tests (#12)

* V2 client test for whisper API

* Fixed test outputs, refactored tests into integration folder

* Added minor unit tests for callback related functions

* Fixed typo in filename

* Minor pre-commit config fix

Author: chandrasekharan-zipstack

.pre-commit-config.yaml

@@ -17,6 +17,7 @@ repos:
         exclude_types:
           - "markdown"
       - id: end-of-file-fixer
+        exclude: "tests/test_data/.*"
       - id: check-yaml
         args: [--unsafe]
       - id: check-added-large-files
@@ -65,9 +66,7 @@ repos:
         args: [--max-line-length=120]
         exclude: |
           (?x)^(
-              .*migrations/.*\.py|
-              unstract-core/tests/.*|
-              pkgs/unstract-flags/src/unstract/flags/evaluation_.*\.py|
+              tests/test_data/.*|
           )$
   - repo: https://github.com/pycqa/isort
     rev: 5.13.2

src/unstract/llmwhisperer/client_v2.py

@@ -18,16 +18,15 @@
     LLMWhispererClientException: Exception raised for errors in the LLMWhispererClient.
 """
 
+import copy
 import json
 import logging
 import os
-from typing import IO
-import copy
 import time
+from typing import IO
 
 import requests
 
-
 BASE_URL = "https://llmwhisperer-api.unstract.com/api/v2"
 
 
@@ -63,9 +62,7 @@ class LLMWhispererClientV2:
     client's activities and errors.
     """
 
-    formatter = logging.Formatter(
-        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-    )
+    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
     logger = logging.getLogger(__name__)
     log_stream_handler = logging.StreamHandler()
     log_stream_handler.setFormatter(formatter)
@@ -251,12 +248,10 @@ def whisper(
         should_stream = False
         if url == "":
             if stream is not None:
-
                 should_stream = True
 
                 def generate():
-                    for chunk in stream:
-                        yield chunk
+                    yield from stream
 
                 req = requests.Request(
                     "POST",
@@ -302,41 +297,31 @@ def generate():
                     message["extraction"] = {}
                     return message
                 if status["status"] == "processing":
-                    self.logger.debug(
-                        f"Whisper-hash:{whisper_hash} | STATUS: processing..."
-                    )
+                    self.logger.debug(f"Whisper-hash:{whisper_hash} | STATUS: processing...")
                 elif status["status"] == "delivered":
-                    self.logger.debug(
-                        f"Whisper-hash:{whisper_hash} | STATUS: Already delivered!"
-                    )
+                    self.logger.debug(f"Whisper-hash:{whisper_hash} | STATUS: Already delivered!")
                     raise LLMWhispererClientException(
                         {
                             "status_code": -1,
                             "message": "Whisper operation already delivered",
                         }
                     )
                 elif status["status"] == "unknown":
-                    self.logger.debug(
-                        f"Whisper-hash:{whisper_hash} | STATUS: unknown..."
-                    )
+                    self.logger.debug(f"Whisper-hash:{whisper_hash} | STATUS: unknown...")
                     raise LLMWhispererClientException(
                         {
                             "status_code": -1,
                             "message": "Whisper operation status unknown",
                         }
                     )
                 elif status["status"] == "failed":
-                    self.logger.debug(
-                        f"Whisper-hash:{whisper_hash} | STATUS: failed..."
-                    )
+                    self.logger.debug(f"Whisper-hash:{whisper_hash} | STATUS: failed...")
                     message["status_code"] = -1
                     message["message"] = "Whisper operation failed"
                     message["extraction"] = {}
                     return message
                 elif status["status"] == "processed":
-                    self.logger.debug(
-                        f"Whisper-hash:{whisper_hash} | STATUS: processed!"
-                    )
+                    self.logger.debug(f"Whisper-hash:{whisper_hash} | STATUS: processed!")
                     resultx = self.whisper_retrieve(whisper_hash=whisper_hash)
                     if resultx["status_code"] == 200:
                         message["status_code"] = 200
@@ -451,7 +436,6 @@ def register_webhook(self, url: str, auth_token: str, webhook_name: str) -> dict
         Raises:
             LLMWhispererClientException: If the API request fails, it raises an exception with
                                             the error message and status code returned by the API.
-
         """
 
         data = {
@@ -489,7 +473,6 @@ def get_webhook_details(self, webhook_name: str) -> dict:
         Raises:
             LLMWhispererClientException: If the API request fails, it raises an exception with
                                             the error message and status code returned by the API.
-
         """
 
         url = f"{self.base_url}/whisper-manage-callback"

tests/client_test_v2.py

@@ -3,15 +3,21 @@
 import pytest
 
 from unstract.llmwhisperer.client import LLMWhispererClient
+from unstract.llmwhisperer.client_v2 import LLMWhispererClientV2
 
 
 @pytest.fixture(name="client")
 def llm_whisperer_client():
-    # Create an instance of the client
     client = LLMWhispererClient()
     return client
 
 
+@pytest.fixture(name="client_v2")
+def llm_whisperer_client_v2():
+    client = LLMWhispererClientV2()
+    return client
+
+
 @pytest.fixture(name="data_dir", scope="session")
 def test_data_dir():
     return os.path.join(os.path.dirname(__file__), "test_data")

tests/conftest.py

@@ -23,9 +23,7 @@ def test_get_usage_info(client):
         "subscription_plan",
         "today_page_count",
     ]
-    assert set(usage_info.keys()) == set(
-        expected_keys
-    ), f"usage_info {usage_info} does not contain the expected keys"
+    assert set(usage_info.keys()) == set(expected_keys), f"usage_info {usage_info} does not contain the expected keys"
 
 
 @pytest.mark.parametrize(
@@ -78,9 +76,7 @@ def test_whisper(self):
     # @unittest.skip("Skipping test_whisper")
     def test_whisper_stream(self):
         client = LLMWhispererClient()
-        download_url = (
-            "https://storage.googleapis.com/pandora-static/samples/bill.jpg.pdf"
-        )
+        download_url = "https://storage.googleapis.com/pandora-static/samples/bill.jpg.pdf"
         # Create a stream of download_url and pass it to whisper
         response_download = requests.get(download_url, stream=True)
         response_download.raise_for_status()
@@ -95,18 +91,14 @@ def test_whisper_stream(self):
     @unittest.skip("Skipping test_whisper_status")
     def test_whisper_status(self):
         client = LLMWhispererClient()
-        response = client.whisper_status(
-            whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a"
-        )
+        response = client.whisper_status(whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a")
         logger.info(response)
         self.assertIsInstance(response, dict)
 
     @unittest.skip("Skipping test_whisper_retrieve")
     def test_whisper_retrieve(self):
         client = LLMWhispererClient()
-        response = client.whisper_retrieve(
-            whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a"
-        )
+        response = client.whisper_retrieve(whisper_hash="7cfa5cbb|5f1d285a7cf18d203de7af1a1abb0a3a")
         logger.info(response)
         self.assertIsInstance(response, dict)
 

tests/integration/__init__.py

@@ -0,0 +1,71 @@
+import logging
+import os
+from difflib import SequenceMatcher, unified_diff
+from pathlib import Path
+
+import pytest
+
+logger = logging.getLogger(__name__)
+
+
+def test_get_usage_info(client_v2):
+    usage_info = client_v2.get_usage_info()
+    logger.info(usage_info)
+    assert isinstance(usage_info, dict), "usage_info should be a dictionary"
+    expected_keys = [
+        "current_page_count",
+        "current_page_count_low_cost",
+        "current_page_count_form",
+        "current_page_count_high_quality",
+        "current_page_count_native_text",
+        "daily_quota",
+        "monthly_quota",
+        "overage_page_count",
+        "subscription_plan",
+        "today_page_count",
+    ]
+    assert set(usage_info.keys()) == set(expected_keys), f"usage_info {usage_info} does not contain the expected keys"
+
+
+@pytest.mark.parametrize(
+    "output_mode, mode, input_file",
+    [
+        ("layout_preserving", "native_text", "credit_card.pdf"),
+        ("layout_preserving", "low_cost", "credit_card.pdf"),
+        ("layout_preserving", "high_quality", "restaurant_invoice_photo.pdf"),
+        ("layout_preserving", "form", "handwritten-form.pdf"),
+        ("text", "native_text", "credit_card.pdf"),
+        ("text", "low_cost", "credit_card.pdf"),
+        ("text", "high_quality", "restaurant_invoice_photo.pdf"),
+        ("text", "form", "handwritten-form.pdf"),
+    ],
+)
+def test_whisper_v2(client_v2, data_dir, output_mode, mode, input_file):
+    file_path = os.path.join(data_dir, input_file)
+    whisper_result = client_v2.whisper(
+        mode=mode, output_mode=output_mode, file_path=file_path, wait_for_completion=True
+    )
+    logger.debug(f"Result for '{output_mode}', '{mode}', " f"'{input_file}: {whisper_result}")
+
+    exp_basename = f"{Path(input_file).stem}.{mode}.{output_mode}.txt"
+    exp_file = os.path.join(data_dir, "expected", exp_basename)
+    with open(exp_file, encoding="utf-8") as f:
+        exp = f.read()
+
+    assert isinstance(whisper_result, dict)
+    assert whisper_result["status_code"] == 200
+
+    # For text based processing, perform a strict match
+    if mode == "native_text" and output_mode == "text":
+        assert whisper_result["extraction"]["result_text"] == exp
+    # For OCR based processing, perform a fuzzy match
+    else:
+        extracted_text = whisper_result["extraction"]["result_text"]
+        similarity = SequenceMatcher(None, extracted_text, exp).ratio()
+        threshold = 0.97
+
+        if similarity < threshold:
+            diff = "\n".join(
+                unified_diff(exp.splitlines(), extracted_text.splitlines(), fromfile="Expected", tofile="Extracted")
+            )
+            pytest.fail(f"Texts are not similar enough: {similarity * 100:.2f}% similarity. Diff:\n{diff}")