tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-24 00:28:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

OCR test fixes

Browse Source
This commit is contained in:
Ishaan Jaffer
2026-01-11 08:00:31 -08:00
parent 080982e157
commit 7aee1fcb81
1 changed files with 28 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/ocr_tests/base_ocr_unit_tests.py
+28 -28
View File
@@ -27,28 +27,6 @@ class BaseOCRTest(ABC):
"""Must return the base OCR call args for the specific provider"""
pass
@pytest.fixture(autouse=True)
def _handle_rate_limits(self):
"""Fixture to handle rate limit errors for all test methods"""
try:
yield
except litellm.RateLimitError as e:
# Check if it's a quota exceeded error
error_msg = str(e)
if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
pytest.skip(f"Quota exceeded - {error_msg}")
else:
pytest.skip(f"Rate limit exceeded - {error_msg}")
except litellm.InternalServerError:
pytest.skip("Model is overloaded")
except litellm.BadRequestError as e:
# Handle URL rejection errors from Vertex AI
error_msg = str(e)
if "URL_REJECTED" in error_msg or "Cannot fetch content from the provided URL" in error_msg:
pytest.skip(f"URL rejected by provider - {error_msg}")
else:
raise
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_basic_ocr_with_url(self, sync_mode):
@@ -124,9 +102,20 @@ class BaseOCRTest(ABC):
assert response_cost > 0, "Response cost should be greater than 0"
print("response_cost=", response_cost)
except (litellm.RateLimitError, litellm.InternalServerError):
# Re-raise these errors so the fixture can handle them
raise
except litellm.RateLimitError as e:
error_msg = str(e)
if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
pytest.skip(f"Quota exceeded - {error_msg}")
else:
pytest.skip(f"Rate limit exceeded - {error_msg}")
except litellm.InternalServerError:
pytest.skip("Model is overloaded")
except litellm.BadRequestError as e:
error_msg = str(e)
if "URL_REJECTED" in error_msg or "Cannot fetch content from the provided URL" in error_msg:
pytest.skip(f"URL rejected by provider - {error_msg}")
else:
pytest.fail(f"OCR call failed: {str(e)}")
except Exception as e:
pytest.fail(f"OCR call failed: {str(e)}")
@@ -170,9 +159,20 @@ class BaseOCRTest(ABC):
print(f" - pages_processed: {response.usage_info.pages_processed}")
print(f" - doc_size_bytes: {response.usage_info.doc_size_bytes}")
except (litellm.RateLimitError, litellm.InternalServerError):
# Re-raise these errors so the fixture can handle them
raise
except litellm.RateLimitError as e:
error_msg = str(e)
if "Quota exceeded" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
pytest.skip(f"Quota exceeded - {error_msg}")
else:
pytest.skip(f"Rate limit exceeded - {error_msg}")
except litellm.InternalServerError:
pytest.skip("Model is overloaded")
except litellm.BadRequestError as e:
error_msg = str(e)
if "URL_REJECTED" in error_msg or "Cannot fetch content from the provided URL" in error_msg:
pytest.skip(f"URL rejected by provider - {error_msg}")
else:
pytest.fail(f"OCR response structure test failed: {str(e)}")
except Exception as e:
pytest.fail(f"OCR response structure test failed: {str(e)}")