LabelingSystem/backend/test/test_export_api.py

"""
Tests for Export API endpoints.
Tests data export functionality in various formats.
"""
import pytest
import json
import os
import uuid
from fastapi.testclient import TestClient
from main import app
from database import get_db_connection, init_database
from test.auth_test_helper import create_test_token


# Test client
client = TestClient(app)


# Test data
TEST_PROJECT = {
    "name": "Export Test Project",
    "description": "Project for testing export functionality",
    "config": "<View><Image name='image' value='$image'/><Choices name='choice' toName='image'><Choice value='Cat'/><Choice value='Dog'/></Choices></View>"
}


@pytest.fixture(scope="module")
def setup_database():
    """Initialize database before tests."""
    init_database()
    yield


@pytest.fixture(scope="module")
def admin_token(setup_database):
    """Create admin user via token cache and get token."""
    user_id = str(uuid.uuid4())
    username = f"export_admin_{uuid.uuid4().hex[:8]}"
    email = f"{username}@test.com"

    # Ensure user exists in local DB
    with get_db_connection() as conn:
        cursor = conn.cursor()
        cursor.execute(
            "INSERT OR IGNORE INTO users (id, username, email, role) VALUES (?, ?, ?, ?)",
            (user_id, username, email, "admin"),
        )

    return create_test_token({
        "id": user_id,
        "username": username,
        "email": email,
        "role": "admin",
    })


@pytest.fixture(scope="module")
def test_project(admin_token):
    """Create a test project with tasks and annotations."""
    headers = {"Authorization": f"Bearer {admin_token}"}
    
    # Create project
    response = client.post("/api/projects", json=TEST_PROJECT, headers=headers)
    assert response.status_code == 201
    project = response.json()
    project_id = project["id"]
    
    # Create tasks
    task_ids = []
    for i in range(3):
        task_data = {
            "project_id": project_id,
            "name": f"Test Task {i+1}",
            "data": {"image": f"https://example.com/image{i+1}.jpg"}
        }
        response = client.post("/api/tasks", json=task_data, headers=headers)
        assert response.status_code == 201
        task_ids.append(response.json()["id"])
    
    # Create annotations for first two tasks
    for i, task_id in enumerate(task_ids[:2]):
        annotation_data = {
            "task_id": task_id,
            "user_id": "test_user",
            "result": {
                "annotations": [
                    {
                        "id": f"ann_{i}",
                        "type": "rectanglelabels",
                        "value": {
                            "x": 10 + i * 5,
                            "y": 20 + i * 5,
                            "width": 30,
                            "height": 40,
                            "rectanglelabels": ["Cat" if i == 0 else "Dog"]
                        }
                    }
                ]
            }
        }
        response = client.post("/api/annotations", json=annotation_data, headers=headers)
        assert response.status_code == 201
    
    # Update first task to completed
    response = client.put(
        f"/api/tasks/{task_ids[0]}",
        json={"status": "completed"},
        headers=headers
    )
    
    yield {
        "project_id": project_id,
        "task_ids": task_ids
    }
    
    # Cleanup: delete project (cascades to tasks and annotations)
    client.delete(f"/api/projects/{project_id}", headers=headers)


class TestExportAPI:
    """Test cases for Export API."""
    
    def test_export_json_format(self, admin_token, test_project):
        """Test exporting data in JSON format."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "json",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["project_id"] == project_id
        assert data["format"] == "json"
        assert data["status"] == "completed"
        assert data["total_tasks"] == 3
        assert data["download_url"] is not None
    
    def test_export_csv_format(self, admin_token, test_project):
        """Test exporting data in CSV format."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "csv",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["format"] == "csv"
        assert data["status"] == "completed"
    
    def test_export_coco_format(self, admin_token, test_project):
        """Test exporting data in COCO format."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "coco",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["format"] == "coco"
        assert data["status"] == "completed"
    
    def test_export_yolo_format(self, admin_token, test_project):
        """Test exporting data in YOLO format."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "yolo",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["format"] == "yolo"
        assert data["status"] == "completed"
    
    def test_export_with_status_filter(self, admin_token, test_project):
        """Test exporting with status filter."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        # Export only completed tasks
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "json",
                "status_filter": "completed",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["status_filter"] == "completed"
        # Should have fewer tasks than total
        assert data["total_tasks"] <= 3
    
    def test_export_nonexistent_project(self, admin_token):
        """Test exporting from non-existent project."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        
        response = client.post(
            "/api/projects/nonexistent_project/export",
            json={
                "format": "json",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 404
    
    def test_get_export_status(self, admin_token, test_project):
        """Test getting export job status."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        # First create an export
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "json",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        export_id = response.json()["id"]
        
        # Get status
        response = client.get(f"/api/exports/{export_id}/status", headers=headers)
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["id"] == export_id
        assert data["status"] == "completed"
        assert data["progress"] == 1.0
    
    def test_download_export(self, admin_token, test_project):
        """Test downloading exported file."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        # First create an export
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "json",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        export_id = response.json()["id"]
        
        # Download
        response = client.get(f"/api/exports/{export_id}/download", headers=headers)
        
        assert response.status_code == 200
        assert response.headers["content-type"] == "application/json"
        
        # Verify content is valid JSON
        content = response.json()
        assert "project_id" in content
        assert "tasks" in content
    
    def test_get_export_job_details(self, admin_token, test_project):
        """Test getting export job details."""
        headers = {"Authorization": f"Bearer {admin_token}"}
        project_id = test_project["project_id"]
        
        # First create an export
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "json",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 200
        export_id = response.json()["id"]
        
        # Get details
        response = client.get(f"/api/exports/{export_id}", headers=headers)
        
        assert response.status_code == 200
        data = response.json()
        
        assert data["id"] == export_id
        assert data["project_id"] == project_id
        assert data["format"] == "json"
        assert data["status"] == "completed"


class TestExportPermissions:
    """Test cases for export permissions."""
    
    def test_non_admin_cannot_export(self, setup_database, test_project):
        """Test that non-admin users cannot export data."""
        user_id = str(uuid.uuid4())
        username = f"regular_user_{uuid.uuid4().hex[:8]}"
        email = f"{username}@test.com"

        # Ensure user exists in local DB as annotator
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(
                "INSERT OR IGNORE INTO users (id, username, email, role) VALUES (?, ?, ?, ?)",
                (user_id, username, email, "annotator"),
            )

        token = create_test_token({
            "id": user_id,
            "username": username,
            "email": email,
            "role": "annotator",
        })
        
        # Try to export
        headers = {"Authorization": f"Bearer {token}"}
        project_id = test_project["project_id"]
        
        response = client.post(
            f"/api/projects/{project_id}/export",
            json={
                "format": "json",
                "status_filter": "all",
                "include_metadata": True
            },
            headers=headers
        )
        
        assert response.status_code == 403