wehub-resource-sync
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
class _FakeBlobServiceClient:
|
|
@classmethod
|
|
def from_connection_string(cls, connection_string):
|
|
return cls()
|
|
|
|
def get_container_client(self, container_name):
|
|
return None
|
|
|
|
|
|
azure_module = types.ModuleType("azure")
|
|
azure_storage_module = types.ModuleType("azure.storage")
|
|
azure_blob_module = types.ModuleType("azure.storage.blob")
|
|
azure_blob_module.BlobServiceClient = _FakeBlobServiceClient
|
|
sys.modules.setdefault("azure", azure_module)
|
|
sys.modules.setdefault("azure.storage", azure_storage_module)
|
|
sys.modules.setdefault("azure.storage.blob", azure_blob_module)
|
|
|
|
module_path = Path(__file__).resolve().parents[1] / "gpt_researcher" / "document" / "azure_document_loader.py"
|
|
spec = importlib.util.spec_from_file_location("azure_document_loader", module_path)
|
|
azure_document_loader = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(azure_document_loader)
|
|
AzureDocumentLoader = azure_document_loader.AzureDocumentLoader
|
|
|
|
|
|
class _Blob:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
|
|
class _BlobDownload:
|
|
def __init__(self, content):
|
|
self._content = content
|
|
|
|
def readall(self):
|
|
return self._content
|
|
|
|
|
|
class _BlobClient:
|
|
def __init__(self, content):
|
|
self._content = content
|
|
|
|
def download_blob(self):
|
|
return _BlobDownload(self._content)
|
|
|
|
|
|
class _Container:
|
|
def __init__(self, blobs):
|
|
self._blobs = blobs
|
|
|
|
def list_blobs(self):
|
|
return [_Blob(name) for name in self._blobs]
|
|
|
|
def get_blob_client(self, name):
|
|
return _BlobClient(self._blobs[name])
|
|
|
|
|
|
class TestAzureDocumentLoader(unittest.IsolatedAsyncioTestCase):
|
|
async def test_load_creates_parent_directories_for_virtual_blob_paths(self):
|
|
loader = AzureDocumentLoader.__new__(AzureDocumentLoader)
|
|
loader.container = _Container({"reports/2026/summary.txt": b"hello"})
|
|
|
|
file_paths = await loader.load()
|
|
|
|
self.assertEqual(len(file_paths), 1)
|
|
downloaded_file = Path(file_paths[0])
|
|
self.assertEqual(downloaded_file.name, "summary.txt")
|
|
self.assertEqual(downloaded_file.read_bytes(), b"hello")
|
|
self.assertIn("reports", downloaded_file.parts)
|
|
self.assertIn("2026", downloaded_file.parts)
|
|
|
|
async def test_load_rejects_blob_names_that_escape_temp_dir(self):
|
|
loader = AzureDocumentLoader.__new__(AzureDocumentLoader)
|
|
loader.container = _Container({"../escape.txt": b"nope"})
|
|
|
|
with self.assertRaises(ValueError):
|
|
await loader.load()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|