LabelingSystem/backend/test/test_annotation_api.py

"""
Unit tests for Annotation API endpoints.
Tests CRUD operations for annotations.
"""
import pytest
import os
import json
from fastapi.testclient import TestClient

# Use a test database
TEST_DB_PATH = "test_annotation_annotation_platform.db"


@pytest.fixture(scope="function", autouse=True)
def setup_test_db():
    """Setup test database before each test and cleanup after."""
    # Set test database path
    original_db_path = os.environ.get("DATABASE_PATH")
    os.environ["DATABASE_PATH"] = TEST_DB_PATH
    
    # Remove existing test database
    if os.path.exists(TEST_DB_PATH):
        os.remove(TEST_DB_PATH)
    
    # Import after setting env var
    from database import init_database
    init_database()
    
    yield
    
    # Cleanup
    if os.path.exists(TEST_DB_PATH):
        os.remove(TEST_DB_PATH)
    
    # Restore original path
    if original_db_path:
        os.environ["DATABASE_PATH"] = original_db_path
    elif "DATABASE_PATH" in os.environ:
        del os.environ["DATABASE_PATH"]


@pytest.fixture(scope="function")
def test_client():
    """Create a test client."""
    from main import app
    return TestClient(app)


@pytest.fixture(scope="function")
def sample_project(test_client):
    """Create a sample project for testing."""
    project_data = {
        "name": "Test Project",
        "description": "Test Description",
        "config": "<View><Image name='img' value='$image'/></View>"
    }
    response = test_client.post("/api/projects", json=project_data)
    return response.json()


@pytest.fixture(scope="function")
def sample_task(test_client, sample_project):
    """Create a sample task for testing."""
    task_data = {
        "project_id": sample_project["id"],
        "name": "Test Task",
        "data": {"image_url": "https://example.com/image.jpg"},
        "assigned_to": "user_001"
    }
    response = test_client.post("/api/tasks", json=task_data)
    return response.json()


def test_list_annotations_empty(test_client):
    """Test listing annotations when database is empty."""
    response = test_client.get("/api/annotations")
    assert response.status_code == 200
    assert response.json() == []


def test_create_annotation(test_client, sample_task):
    """Test creating a new annotation."""
    annotation_data = {
        "task_id": sample_task["id"],
        "user_id": "user_001",
        "result": {
            "annotations": [
                {
                    "id": "ann_1",
                    "type": "rectanglelabels",
                    "value": {
                        "x": 10,
                        "y": 20,
                        "width": 100,
                        "height": 50,
                        "rectanglelabels": ["Cat"]
                    }
                }
            ]
        }
    }
    
    response = test_client.post("/api/annotations", json=annotation_data)
    assert response.status_code == 201
    
    data = response.json()
    assert data["task_id"] == annotation_data["task_id"]
    assert data["user_id"] == annotation_data["user_id"]
    assert data["result"] == annotation_data["result"]
    assert "id" in data
    assert data["id"].startswith("ann_")
    assert "created_at" in data
    assert "updated_at" in data


def test_create_annotation_invalid_task(test_client):
    """Test creating an annotation with invalid task_id fails."""
    annotation_data = {
        "task_id": "nonexistent_task",
        "user_id": "user_001",
        "result": {"annotations": []}
    }
    
    response = test_client.post("/api/annotations", json=annotation_data)
    assert response.status_code == 404
    assert "not found" in response.json()["detail"].lower()


def test_get_annotation(test_client, sample_task):
    """Test getting an annotation by ID."""
    # Create an annotation first
    annotation_data = {
        "task_id": sample_task["id"],
        "user_id": "user_001",
        "result": {"annotations": [{"id": "ann_1", "value": "test"}]}
    }
    create_response = test_client.post("/api/annotations", json=annotation_data)
    annotation_id = create_response.json()["id"]
    
    # Get the annotation
    response = test_client.get(f"/api/annotations/{annotation_id}")
    assert response.status_code == 200
    
    data = response.json()
    assert data["id"] == annotation_id
    assert data["task_id"] == annotation_data["task_id"]
    assert data["user_id"] == annotation_data["user_id"]


def test_get_annotation_not_found(test_client):
    """Test getting a non-existent annotation returns 404."""
    response = test_client.get("/api/annotations/nonexistent_id")
    assert response.status_code == 404
    assert "not found" in response.json()["detail"].lower()


def test_update_annotation(test_client, sample_task):
    """Test updating an annotation."""
    # Create an annotation first
    annotation_data = {
        "task_id": sample_task["id"],
        "user_id": "user_001",
        "result": {"annotations": [{"id": "ann_1", "value": "original"}]}
    }
    create_response = test_client.post("/api/annotations", json=annotation_data)
    annotation_id = create_response.json()["id"]
    
    # Update the annotation
    update_data = {
        "result": {"annotations": [{"id": "ann_1", "value": "updated"}]}
    }
    response = test_client.put(f"/api/annotations/{annotation_id}", json=update_data)
    assert response.status_code == 200
    
    data = response.json()
    assert data["result"] == update_data["result"]
    assert data["task_id"] == annotation_data["task_id"]  # Task ID unchanged
    assert data["user_id"] == annotation_data["user_id"]  # User ID unchanged


def test_update_annotation_not_found(test_client):
    """Test updating a non-existent annotation returns 404."""
    update_data = {"result": {"annotations": []}}
    response = test_client.put("/api/annotations/nonexistent_id", json=update_data)
    assert response.status_code == 404


def test_list_annotations_after_creation(test_client, sample_task):
    """Test listing annotations after creating some."""
    # Create multiple annotations
    for i in range(3):
        annotation_data = {
            "task_id": sample_task["id"],
            "user_id": f"user_{i:03d}",
            "result": {"annotations": [{"id": f"ann_{i}", "value": f"test_{i}"}]}
        }
        test_client.post("/api/annotations", json=annotation_data)
    
    # List annotations
    response = test_client.get("/api/annotations")
    assert response.status_code == 200
    
    data = response.json()
    assert len(data) == 3
    assert all("id" in annotation for annotation in data)
    assert all("created_at" in annotation for annotation in data)
    assert all("updated_at" in annotation for annotation in data)


def test_list_annotations_filter_by_task(test_client, sample_project):
    """Test filtering annotations by task_id."""
    # Create two tasks
    task1_data = {
        "project_id": sample_project["id"],
        "name": "Task 1",
        "data": {"image_url": "https://example.com/image1.jpg"}
    }
    task1_response = test_client.post("/api/tasks", json=task1_data)
    task1 = task1_response.json()
    
    task2_data = {
        "project_id": sample_project["id"],
        "name": "Task 2",
        "data": {"image_url": "https://example.com/image2.jpg"}
    }
    task2_response = test_client.post("/api/tasks", json=task2_data)
    task2 = task2_response.json()
    
    # Create annotations for both tasks
    ann1_data = {
        "task_id": task1["id"],
        "user_id": "user_001",
        "result": {"annotations": [{"id": "ann_1"}]}
    }
    test_client.post("/api/annotations", json=ann1_data)
    
    ann2_data = {
        "task_id": task2["id"],
        "user_id": "user_001",
        "result": {"annotations": [{"id": "ann_2"}]}
    }
    test_client.post("/api/annotations", json=ann2_data)
    
    # Filter by first task
    response = test_client.get(f"/api/annotations?task_id={task1['id']}")
    assert response.status_code == 200
    
    data = response.json()
    assert len(data) == 1
    assert data[0]["task_id"] == task1["id"]


def test_list_annotations_filter_by_user(test_client, sample_task):
    """Test filtering annotations by user_id."""
    # Create annotations for different users
    ann1_data = {
        "task_id": sample_task["id"],
        "user_id": "user_001",
        "result": {"annotations": [{"id": "ann_1"}]}
    }
    test_client.post("/api/annotations", json=ann1_data)
    
    ann2_data = {
        "task_id": sample_task["id"],
        "user_id": "user_002",
        "result": {"annotations": [{"id": "ann_2"}]}
    }
    test_client.post("/api/annotations", json=ann2_data)
    
    # Filter by first user
    response = test_client.get("/api/annotations?user_id=user_001")
    assert response.status_code == 200
    
    data = response.json()
    assert len(data) == 1
    assert data[0]["user_id"] == "user_001"


def test_get_task_annotations(test_client, sample_task):
    """Test getting all annotations for a specific task."""
    # Create annotations for the task
    for i in range(2):
        annotation_data = {
            "task_id": sample_task["id"],
            "user_id": f"user_{i:03d}",
            "result": {"annotations": [{"id": f"ann_{i}"}]}
        }
        test_client.post("/api/annotations", json=annotation_data)
    
    # Get task annotations using the alternative endpoint
    response = test_client.get(f"/api/annotations/tasks/{sample_task['id']}/annotations")
    assert response.status_code == 200
    
    data = response.json()
    assert len(data) == 2
    assert all(annotation["task_id"] == sample_task["id"] for annotation in data)


def test_get_task_annotations_not_found(test_client):
    """Test getting annotations for non-existent task returns 404."""
    response = test_client.get("/api/annotations/tasks/nonexistent_id/annotations")
    assert response.status_code == 404


def test_annotation_json_serialization(test_client, sample_task):
    """Test that complex JSON data is properly serialized and deserialized."""
    complex_result = {
        "annotations": [
            {
                "id": "ann_1",
                "type": "rectanglelabels",
                "value": {
                    "x": 10.5,
                    "y": 20.3,
                    "width": 100,
                    "height": 50,
                    "rectanglelabels": ["Cat", "Animal"]
                }
            },
            {
                "id": "ann_2",
                "type": "choices",
                "value": {
                    "choices": ["Option A"]
                }
            }
        ],
        "metadata": {
            "duration": 120,
            "quality": "high"
        }
    }
    
    annotation_data = {
        "task_id": sample_task["id"],
        "user_id": "user_001",
        "result": complex_result
    }
    
    # Create annotation
    create_response = test_client.post("/api/annotations", json=annotation_data)
    assert create_response.status_code == 201
    annotation_id = create_response.json()["id"]
    
    # Get annotation and verify data integrity
    get_response = test_client.get(f"/api/annotations/{annotation_id}")
    assert get_response.status_code == 200
    
    data = get_response.json()
    assert data["result"] == complex_result
    assert data["result"]["annotations"][0]["value"]["x"] == 10.5
    assert data["result"]["metadata"]["duration"] == 120