add testing to splitter endpoints

Author: laurasgkadri98

app/fastapi_utils.py

@@ -77,8 +77,8 @@ def get_parameters_info(params):
         if param.annotation == bytes:
             param_info = ParameterInfo(name=param.name, type="bytes")
             parameters_info.append(param_info)
-        elif hasattr(param.annotation, "schema"):
-            schema = param.annotation.schema()
+        elif hasattr(param.annotation, "model_json_schema"):
+            schema = param.annotation.model_json_schema()
             param_info = extract_fields_from_schema(schema)
             parameters_info.extend(param_info)
         else:
@@ -101,7 +101,7 @@ def get_return_type_info(return_type: Type[BaseModel]):
         A list of ParameterInfo objects representing the return type fields.
 
     """
-    if hasattr(return_type, "schema"):
+    if hasattr(return_type, "model_json_schema"):
         schema = return_type.model_json_schema()
         return extract_fields_from_schema(schema)
     return [ParameterInfo(name="return", type=str(return_type.__name__))]

requirements.txt

@@ -2,5 +2,6 @@ fastapi==0.111.1
 langchain==0.2.11
 pydantic==2.8.2
 pymupdf==1.24.9
+pytest==8.3.2
 python_pptx==0.6.23
 PyYAML==6.0.1

tests/__init__.py

@@ -0,0 +1 @@
+"""Tests module."""

tests/conftest.py

@@ -0,0 +1,11 @@
+import pytest
+from unittest.mock import patch
+
+# Mock API key for testing
+MOCK_API_KEY = "test_api_key"
+
+@pytest.fixture(autouse=True)
+def mock_api_key():
+    """Mock the API key for testing."""
+    with patch('app.config.CONFIG.flowkit_python_api_key', MOCK_API_KEY):
+        yield

tests/test_endpoints_splitter.py

@@ -0,0 +1,142 @@
+import base64
+from fastapi import HTTPException
+import pytest
+from fastapi.testclient import TestClient
+from app.app import app
+from app.endpoints.splitter import validate_request
+from app.models.splitter import SplitterRequest
+from tests.conftest import MOCK_API_KEY
+
+# Create a test client
+client = TestClient(app)
+
+def encode_file_to_base64(file_path):
+    """Encode a file to base64 string."""
+    with open(file_path, "rb") as file:
+        return base64.b64encode(file.read()).decode('utf-8')
+
+@pytest.mark.asyncio
+async def test_split_ppt():
+    """Test splitting text in a PowerPoint document into chunks."""
+    ppt_content_base64 = encode_file_to_base64("./tests/test_files/test_presentation.pptx")
+    request_payload = {
+        "document_content": ppt_content_base64,
+        "chunk_size": 100,
+        "chunk_overlap": 10
+    }
+    response = client.post(
+        "/splitter/ppt",
+        json=request_payload,
+        headers={"api-key": MOCK_API_KEY}
+    )
+    if response.status_code != 200:
+        print(f"Response status code: {response.status_code}")
+        print(f"Response content: {response.json()}")
+    assert response.status_code == 200
+    assert "chunks" in response.json()
+
+@pytest.mark.asyncio
+async def test_split_py():
+    """Test splitting Python code into chunks."""
+    python_code = """
+    def hello_world():
+        print("Hello, world!")
+    """
+    python_code_base64 = base64.b64encode(python_code.encode()).decode('utf-8')
+    request_payload = {
+        "document_content": python_code_base64,
+        "chunk_size": 50,
+        "chunk_overlap": 5
+    }
+    response = client.post(
+        "/splitter/py",
+        json=request_payload,
+        headers={"api-key": MOCK_API_KEY}
+    )
+    assert response.status_code == 200
+    assert "chunks" in response.json()
+
+@pytest.mark.asyncio
+async def test_split_pdf():
+    """Test splitting text in a PDF document into chunks."""
+    pdf_content_base64 = encode_file_to_base64("./tests/test_files/test_document.pdf")
+    request_payload = {
+        "document_content": pdf_content_base64,
+        "chunk_size": 200,
+        "chunk_overlap": 20
+    }
+    response = client.post(
+        "/splitter/pdf",
+        json=request_payload,
+        headers={"api-key": MOCK_API_KEY}
+    )
+    assert response.status_code == 200
+    assert "chunks" in response.json()
+
+# Define test cases for validate_request()
+validate_request_test_cases = [
+    # Test case 1: valid request
+    (
+        SplitterRequest(
+            document_content="dGVzdA==",  # base64 for "test"
+            chunk_size=100,
+            chunk_overlap=10
+        ),
+        MOCK_API_KEY,
+        None,
+    ),
+    # Test case: invalid API key
+    (
+        SplitterRequest(
+            document_content="dGVzdA==",
+            chunk_size=100,
+            chunk_overlap=10
+        ),
+        "invalid_api_key",
+        HTTPException(status_code=401, detail="Invalid API key"),
+    ),
+    # Test case 2: missing document content
+    (
+        SplitterRequest(
+            document_content="",
+            chunk_size=100,
+            chunk_overlap=10
+        ),
+        MOCK_API_KEY,
+        HTTPException(status_code=400, detail="No document content provided"),
+    ),
+    # Test case 4: invalid chunk size
+    (
+        SplitterRequest(
+            document_content="dGVzdA==",
+            chunk_size=0,
+            chunk_overlap=10
+        ),
+        MOCK_API_KEY,
+        HTTPException(status_code=400, detail="No chunk size provided"),
+    ),
+    # Test case 5: invalid chunk overlap
+    (
+        SplitterRequest(
+            document_content="dGVzdA==",
+            chunk_size=100,
+            chunk_overlap=-1
+        ),
+        MOCK_API_KEY,
+        HTTPException(status_code=400, detail="Chunk overlap must be greater than or equal to 0"),
+    ),
+]
+
+@pytest.mark.parametrize("api_request, api_key, expected_exception", validate_request_test_cases)
+def test_validate_request(api_request, api_key, expected_exception):
+    """Test the validate_request function with various scenarios."""
+    if expected_exception:
+        with pytest.raises(HTTPException) as exc_info:
+            validate_request(api_request, api_key)
+        assert exc_info.value.status_code == expected_exception.status_code
+        assert exc_info.value.detail == expected_exception.detail
+    else:
+        try:
+            validate_request(api_request, api_key)
+        except HTTPException:
+            pytest.fail("validate_request() raised HTTPException unexpectedly!")