From 47dd35390374acc458480c6bb00f3f7a6b81910e Mon Sep 17 00:00:00 2001
From: Kye Gomez <swarms_wd@Kyes-MacBook-Pro.local>
Date: Fri, 16 Aug 2024 13:38:42 -0400
Subject: [PATCH] [TESTS][API]

---
 .env.examples     |  3 +-
 server/api.py     | 11 +++---
 tests/test_api.py | 97 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 105 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_api.py

diff --git a/.env.examples b/.env.examples
index 209f0c1..9d8ad7c 100644
--- a/.env.examples
+++ b/.env.examples
@@ -1 +1,2 @@
-PINECONE_API_KEYS="your_pinecone_api_key"
\ No newline at end of file
+PINECONE_API_KEYS="your_pinecone_api_key"
+BASE_SWARMS_MEMORY_URL="http:"
\ No newline at end of file
diff --git a/server/api.py b/server/api.py
index f093fce..ef89294 100644
--- a/server/api.py
+++ b/server/api.py
@@ -1,11 +1,12 @@
-from fastapi import FastAPI, HTTPException, Path, Body
-from pydantic import BaseModel, Field
+import os
+import uuid
 from typing import List, Optional
-from loguru import logger
+
 import chromadb
-import uuid
-import os
 import uvicorn
+from fastapi import Body, FastAPI, HTTPException, Path
+from loguru import logger
+from pydantic import BaseModel, Field
 
 app = FastAPI()
 
diff --git a/tests/test_api.py b/tests/test_api.py
new file mode 100644
index 0000000..ff90047
--- /dev/null
+++ b/tests/test_api.py
@@ -0,0 +1,97 @@
+import pytest
+import requests
+import os
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+BASE_URL = os.getenv("BASE_SWARMS_MEMORY_URL")
+
+
+@pytest.fixture(scope="module")
+def collection_name():
+    """
+    Fixture for creating a unique collection name.
+    """
+    return "my_collection"
+
+
+@pytest.fixture(scope="module")
+def collection_url():
+    """
+    Fixture to return the base collection URL.
+    """
+    return f"{BASE_URL}/collections"
+
+
+@pytest.fixture(scope="module")
+def setup_collection(collection_url, collection_name):
+    """
+    Fixture to create a collection before tests and delete it after.
+    """
+    # Create the collection
+    data = {"name": collection_name}
+    response = requests.post(collection_url, json=data)
+    assert response.status_code == 201
+    collection_id = response.json().get("collection_id")
+
+    yield collection_id
+
+    # Delete the collection after tests
+    requests.delete(f"{collection_url}/{collection_id}")
+
+
+def test_create_collection(setup_collection):
+    """
+    Test the creation of a new collection.
+    """
+    # The setup_collection fixture already creates the collection and returns the ID.
+    assert setup_collection is not None
+
+
+def test_add_documents(collection_url, setup_collection):
+    """
+    Test adding documents to a collection.
+    """
+    url = f"{collection_url}/{setup_collection}/documents"
+    data = {
+        "documents": [
+            "This is a document about pineapples",
+            "This is a document about oranges"
+        ],
+    }
+    response = requests.post(url, json=data)
+    assert response.status_code == 200
+    response_json = response.json()
+    assert "document_ids" in response_json
+
+    # Return the first document ID for use in other tests
+    return response_json["document_ids"][0]
+
+
+def test_query_documents(collection_url, setup_collection):
+    """
+    Test querying documents in a collection.
+    """
+    url = f"{collection_url}/{setup_collection}/documents"
+    data = {
+        "query_texts": ["This is a query document about Hawaii"],
+        "n_results": 2
+    }
+    response = requests.get(url, json=data)
+    assert response.status_code == 200
+    response_json = response.json()
+    assert "results" in response_json
+
+
+def test_delete_document(collection_url, setup_collection, test_add_documents):
+    """
+    Test deleting a document from a collection.
+    """
+    document_id = test_add_documents  # Use the document ID from the add documents test
+    url = f"{collection_url}/{setup_collection}/documents/{document_id}"
+    response = requests.delete(url)
+    assert response.status_code == 200
+    response_json = response.json()
+    assert "deleted" in response_json
\ No newline at end of file