add integration test

Author: blink1073

pyproject.toml

@@ -70,6 +70,7 @@ addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
 xfail_strict = true
 filterwarnings = [
     "error",
+    "module:datetime.datetime.utcnow:DeprecationWarning", # from boto3
     "module:builtin type Swig:DeprecationWarning", # from pymupdf
     "module:builtin type swig:DeprecationWarning", # from pymupdf
 ]

tests/test_client_integration.py

@@ -1,16 +1,18 @@
 import os
+import urllib.request
 from collections.abc import Generator
 
 import numpy as np
 import pytest
+from bson import ObjectId
 
 from pymongo_voyageai import PyMongoVoyageAI
 
 if "VOYAGEAI_API_KEY" not in os.environ:
     pytest.skip("Requires VoyageAI API Key.", allow_module_level=True)
 
 if "S3_BUCKET_NAME" not in os.environ:
-    pytest.skip("Requires VoyageAI API Key.", allow_module_level=True)
+    pytest.skip("Requires S3 Bucket name.", allow_module_level=True)
 
 # mypy: disable_error_code="no-untyped-def"
 
@@ -75,6 +77,24 @@ def test_pdf_pages(client: PyMongoVoyageAI):
     client.delete_by_ids([d["_id"] for d in resp])
 
 
+def test_pdf_pages_storage(client: PyMongoVoyageAI):
+    query = "The consequences of a dictator's peace"
+    url = "https://www.fdrlibrary.org/documents/356632/390886/readingcopy.pdf"
+    storage = client._storage
+    object_name = f"{ObjectId()}.pdf"
+    with urllib.request.urlopen(url) as response:
+        storage.client.upload_fileobj(response, storage.root_location, object_name)
+    url = f"s3://{storage.root_location}/{object_name}"
+    images = client.url_to_images(url)
+    resp = client.add_documents(images)
+    client.wait_for_indexing()
+    data = client.similarity_search(query, extract_images=True)
+    assert len(data[0]["inputs"][0].image.tobytes()) > 0
+    assert len(client.get_by_ids([d["_id"] for d in resp])) == len(resp)
+    client.delete_by_ids([d["_id"] for d in resp])
+    storage.client.delete_object(Bucket=storage.root_location, Key=object_name)
+
+
 @pytest.mark.asyncio
 async def test_image_set_async(client: PyMongoVoyageAI):
     url = "hf://datasets/princeton-nlp/CharXiv/val.parquet"