adding tests for files and vector stores

tests/providers/openai/hooks/test_openai.py

@@ -23,10 +23,20 @@
 
 openai = pytest.importorskip("openai")
 
+from unittest.mock import mock_open
+
 from openai.pagination import SyncCursorPage
-from openai.types import CreateEmbeddingResponse, Embedding
-from openai.types.beta import Assistant, AssistantDeleted, Thread, ThreadDeleted
+from openai.types import CreateEmbeddingResponse, Embedding, FileDeleted, FileObject
+from openai.types.beta import (
+    Assistant,
+    AssistantDeleted,
+    Thread,
+    ThreadDeleted,
+    VectorStore,
+    VectorStoreDeleted,
+)
 from openai.types.beta.threads import Message, Run
+from openai.types.beta.vector_stores import VectorStoreFile, VectorStoreFileBatch, VectorStoreFileDeleted
 from openai.types.chat import ChatCompletion
 
 from airflow.models import Connection
@@ -39,7 +49,12 @@
 MESSAGE_ID = "test_message_abc123"
 RUN_ID = "test_run_abc123"
 MODEL = "gpt-4"
+FILE_ID = "test_file_abc123"
+FILE_NAME = "test_file.pdf"
 METADATA = {"modified": "true", "user": "abc123"}
+VECTOR_STORE_ID = "test_vs_abc123"
+VECTOR_STORE_NAME = "Test Vector Store"
+VECTOR_FILE_STORE_BATCH_ID = "test_vfsb_abc123"
 
 
 @pytest.fixture
@@ -161,6 +176,86 @@ def mock_run_list(mock_run):
     return SyncCursorPage[Run](data=[mock_run])
 
 
+@pytest.fixture
+def mock_file():
+    return FileObject(
+        id=FILE_ID,
+        object="file",
+        bytes=120000,
+        created_at=1677610602,
+        filename=FILE_NAME,
+        purpose="assistants",
+        status="processed",
+    )
+
+
+@pytest.fixture
+def mock_file_list(mock_file):
+    return SyncCursorPage[FileObject](data=[mock_file])
+
+
+@pytest.fixture
+def mock_vector_store():
+    return VectorStore(
+        id=VECTOR_STORE_ID,
+        object="vector_store",
+        created_at=1698107661,
+        usage_bytes=123456,
+        last_active_at=1698107661,
+        name=VECTOR_STORE_NAME,
+        bytes=123456,
+        status="completed",
+        file_counts={"in_progress": 0, "completed": 100, "cancelled": 0, "failed": 0, "total": 100},
+        metadata={},
+        last_used_at=1698107661,
+    )
+
+
+@pytest.fixture
+def mock_vector_store_list(mock_vector_store):
+    return SyncCursorPage[VectorStore](data=[mock_vector_store])
+
+
+@pytest.fixture
+def mock_vector_file_store_batch():
+    return VectorStoreFileBatch(
+        id=VECTOR_FILE_STORE_BATCH_ID,
+        object="vector_store.files_batch",
+        created_at=1699061776,
+        vector_store_id=VECTOR_STORE_ID,
+        status="completed",
+        file_counts={
+            "in_progress": 0,
+            "completed": 3,
+            "failed": 0,
+            "cancelled": 0,
+            "total": 0,
+        },
+    )
+
+
+@pytest.fixture
+def mock_vector_file_store_list():
+    return SyncCursorPage[VectorStoreFile](
+        data=[
+            VectorStoreFile(
+                id="test-file-abc123",
+                object="vector_store.file",
+                created_at=1699061776,
+                vector_store_id=VECTOR_STORE_ID,
+                status="completed",
+            ),
+            VectorStoreFile(
+                id="test-file-abc456",
+                object="vector_store.file",
+                created_at=1699061776,
+                vector_store_id=VECTOR_STORE_ID,
+                status="completed",
+            ),
+        ]
+    )
+
+
 def test_create_chat_completion(mock_openai_hook, mock_completion):
     messages = [
         {"role": "system", "content": "You are a helpful assistant."},
@@ -269,6 +364,14 @@ def test_create_run(mock_openai_hook, mock_run):
     assert run.id == RUN_ID
 
 
+def test_create_run_and_poll(mock_openai_hook, mock_run):
+    thread_id = THREAD_ID
+    assistant_id = ASSISTANT_ID
+    mock_openai_hook.conn.beta.threads.runs.create_and_poll.return_value = mock_run
+    run = mock_openai_hook.create_run_and_poll(thread_id=thread_id, assistant_id=assistant_id)
+    assert run.id == RUN_ID
+
+
 def test_get_runs(mock_openai_hook, mock_run_list):
     mock_openai_hook.conn.beta.threads.runs.list.return_value = mock_run_list
     runs = mock_openai_hook.get_runs(thread_id=THREAD_ID)
@@ -296,6 +399,117 @@ def test_create_embeddings(mock_openai_hook, mock_embeddings_response):
     assert embeddings == [0.1, 0.2, 0.3]
 
 
+@patch("builtins.open", new_callable=mock_open, read_data="test-data")
+def test_upload_file(mock_file_open, mock_openai_hook, mock_file):
+    mock_file.name = FILE_NAME
+    mock_file.purpose = "assistants"
+    mock_openai_hook.conn.files.create.return_value = mock_file
+    file = mock_openai_hook.upload_file(file=mock_file_open(), purpose="assistants")
+    assert file.name == FILE_NAME
+    assert file.purpose == "assistants"
+
+
+def test_get_file(mock_openai_hook, mock_file):
+    mock_openai_hook.conn.files.retrieve.return_value = mock_file
+    file = mock_openai_hook.get_file(file_id=FILE_ID)
+    assert file.id == FILE_ID
+    assert file.filename == FILE_NAME
+
+
+def test_get_files(mock_openai_hook, mock_file_list):
+    mock_openai_hook.conn.files.list.return_value = mock_file_list
+    files = mock_openai_hook.get_files()
+    assert isinstance(files, list)
+
+
+def test_get_file_by_name(mock_openai_hook, mock_file_list):
+    mock_openai_hook.conn.files.list.return_value = mock_file_list
+    file = mock_openai_hook.get_file_by_name(file_name=FILE_NAME)
+    assert file.id == FILE_ID
+    assert file.filename == FILE_NAME
+
+
+def test_delete_file(mock_openai_hook):
+    delete_response = FileDeleted(id=FILE_ID, object="file", deleted=True)
+    mock_openai_hook.conn.files.delete.return_value = delete_response
+    file_deleted = mock_openai_hook.delete_file(file_id=FILE_ID)
+    assert file_deleted.deleted
+
+
+def test_create_vector_store(mock_openai_hook, mock_vector_store):
+    mock_openai_hook.conn.beta.vector_stores.create.return_value = mock_vector_store
+    vector_store = mock_openai_hook.create_vector_store(name=VECTOR_STORE_NAME)
+    assert vector_store.id == VECTOR_STORE_ID
+    assert vector_store.name == VECTOR_STORE_NAME
+
+
+def test_get_vector_store(mock_openai_hook, mock_vector_store):
+    mock_openai_hook.conn.beta.vector_stores.retrieve.return_value = mock_vector_store
+    vector_store = mock_openai_hook.get_vector_store(vector_store_id=VECTOR_STORE_ID)
+    assert vector_store.id == VECTOR_STORE_ID
+    assert vector_store.name == VECTOR_STORE_NAME
+
+
+def test_get_vector_stores(mock_openai_hook, mock_vector_store_list):
+    mock_openai_hook.conn.beta.vector_stores.list.return_value = mock_vector_store_list
+    vector_stores = mock_openai_hook.get_vectors_stores()
+    assert isinstance(vector_stores, list)
+
+
+def test_get_vector_store_by_name(mock_openai_hook, mock_vector_store_list):
+    mock_openai_hook.conn.beta.vector_stores.list.return_value = mock_vector_store_list
+    vector_store = mock_openai_hook.get_vector_store_by_name(vector_store_name=VECTOR_STORE_NAME)
+    assert vector_store.id == VECTOR_STORE_ID
+    assert vector_store.name == VECTOR_STORE_NAME
+
+
+def test_modify_vector_store(mock_openai_hook, mock_vector_store):
+    new_vector_store_name = "New Vector Store"
+    mock_vector_store.name = new_vector_store_name
+    mock_openai_hook.conn.beta.vector_stores.update.return_value = mock_vector_store
+    vector_store = mock_openai_hook.modify_vector_store(
+        vector_store_id=VECTOR_STORE_ID, name=new_vector_store_name
+    )
+    assert vector_store.name == new_vector_store_name
+
+
+def test_delete_vector_store(mock_openai_hook):
+    delete_response = VectorStoreDeleted(id=VECTOR_STORE_ID, object="vector_store.deleted", deleted=True)
+    mock_openai_hook.conn.beta.vector_stores.delete.return_value = delete_response
+    vector_store_deleted = mock_openai_hook.delete_vector_store(vector_store_id=VECTOR_STORE_ID)
+    assert vector_store_deleted.deleted
+
+
+def test_upload_files_to_vector_store(mock_openai_hook, mock_vector_file_store_batch):
+    files = ["file1.txt", "file2.txt", "file3.txt"]
+    mock_openai_hook.conn.beta.vector_stores.file_batches.upload_and_poll.return_value = (
+        mock_vector_file_store_batch
+    )
+    vector_file_store_batch = mock_openai_hook.upload_files_to_vector_store(
+        vector_store_id=VECTOR_STORE_ID, files=files
+    )
+    assert vector_file_store_batch.id == VECTOR_FILE_STORE_BATCH_ID
+    assert vector_file_store_batch.file_counts.completed == len(files)
+
+
+def test_get_vector_store_files(mock_openai_hook, mock_vector_file_store_list):
+    mock_openai_hook.conn.beta.vector_stores.files.list.return_value = mock_vector_file_store_list
+    vector_file_store_list = mock_openai_hook.get_vector_store_files(vector_store_id=VECTOR_STORE_ID)
+    assert isinstance(vector_file_store_list, list)
+
+
+def test_delete_vector_store_file(mock_openai_hook):
+    delete_response = VectorStoreFileDeleted(
+        id="test_file_abc123", object="vector_store.file.deleted", deleted=True
+    )
+    mock_openai_hook.conn.beta.vector_stores.files.delete.return_value = delete_response
+    vector_store_file_deleted = mock_openai_hook.delete_vector_store_file(
+        vector_store_id=VECTOR_STORE_ID, file_id=FILE_ID
+    )
+    assert vector_store_file_deleted.id == FILE_ID
+    assert vector_store_file_deleted.deleted
+
+
 def test_openai_hook_test_connection(mock_openai_hook):
     result, message = mock_openai_hook.test_connection()
     assert result is True