260 lines
8.6 KiB
Python
260 lines
8.6 KiB
Python
import json
|
|
from unittest.mock import AsyncMock, patch
|
|
import pytest
|
|
import httpx
|
|
from starlette.testclient import TestClient
|
|
from server import create_app, outline_client, OUTLINE_API_TOKEN
|
|
|
|
# Helper mock responses
|
|
MOCK_ATTACH_CREATE_SUCCESS = {
|
|
"ok": True,
|
|
"data": {
|
|
"id": "attachment-5678-uuid",
|
|
"name": "test_stream.txt",
|
|
"size": 34,
|
|
"contentType": "text/plain",
|
|
"uploadUrl": "https://s3.example.com/fake-bucket/test_stream.txt?signature=xyz",
|
|
"url": "https://docs.example.com/attachments/attachment-5678-uuid"
|
|
}
|
|
}
|
|
|
|
MOCK_ATTACH_CREATE_ERROR = {
|
|
"ok": False,
|
|
"error": "Document not found"
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_success():
|
|
"""Verify that a successful stream upload coordinates with Outline,
|
|
streams content to S3, and returns the expected metadata.
|
|
"""
|
|
file_content = b"Hello, Outline Wiki Stream upload!"
|
|
content_length = len(file_content)
|
|
document_id = "doc-1234-uuid"
|
|
filename = "test_stream.txt"
|
|
content_type = "text/plain"
|
|
|
|
app = create_app()
|
|
|
|
# Mock OutlineClient attachments.create
|
|
with patch.object(outline_client, "call", new_callable=AsyncMock) as mock_call:
|
|
mock_call.return_value = MOCK_ATTACH_CREATE_SUCCESS
|
|
|
|
# Mock S3 PUT response
|
|
mock_put_response = httpx.Response(
|
|
200, request=httpx.Request("PUT", "https://s3.example.com")
|
|
)
|
|
|
|
# Mock httpx.AsyncClient to intercept put
|
|
with patch("httpx.AsyncClient") as MockClient:
|
|
mock_instance = AsyncMock()
|
|
mock_instance.put.return_value = mock_put_response
|
|
MockClient.return_value.__aenter__.return_value = mock_instance
|
|
|
|
client = TestClient(app)
|
|
response = client.post(
|
|
f"/upload?documentId={document_id}&name={filename}",
|
|
content=file_content,
|
|
headers={
|
|
"Content-Length": str(content_length),
|
|
"Content-Type": content_type,
|
|
"Authorization": f"Bearer {OUTLINE_API_TOKEN}"
|
|
}
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
response_json = response.json()
|
|
assert response_json["ok"] is True
|
|
assert response_json["data"]["id"] == "attachment-5678-uuid"
|
|
|
|
# Verify attachments.create parameters
|
|
mock_call.assert_called_once_with(
|
|
"attachments.create",
|
|
{
|
|
"name": filename,
|
|
"documentId": document_id,
|
|
"size": content_length,
|
|
"contentType": content_type
|
|
}
|
|
)
|
|
|
|
# Verify S3 PUT parameters
|
|
mock_instance.put.assert_called_once()
|
|
args, kwargs = mock_instance.put.call_args
|
|
assert args[0] == "https://s3.example.com/fake-bucket/test_stream.txt?signature=xyz"
|
|
assert kwargs["headers"]["Content-Type"] == content_type
|
|
assert kwargs["headers"]["Content-Length"] == str(content_length)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_unauthorized():
|
|
"""Verify that request with incorrect or missing auth is rejected with 401."""
|
|
app = create_app()
|
|
client = TestClient(app)
|
|
|
|
# Missing Authorization header
|
|
response = client.post(
|
|
"/upload?documentId=doc-123&name=test.txt",
|
|
content=b"content",
|
|
headers={"Content-Length": "7"}
|
|
)
|
|
assert response.status_code == 401
|
|
assert response.json()["error"] == "Unauthorized"
|
|
|
|
# Incorrect Authorization token
|
|
response = client.post(
|
|
"/upload?documentId=doc-123&name=test.txt",
|
|
content=b"content",
|
|
headers={
|
|
"Content-Length": "7",
|
|
"Authorization": "Bearer wrong-token"
|
|
}
|
|
)
|
|
assert response.status_code == 401
|
|
assert response.json()["error"] == "Unauthorized"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_oversized_file():
|
|
"""Verify that oversized files are rejected with 413 Payload Too Large."""
|
|
app = create_app()
|
|
client = TestClient(app)
|
|
|
|
# Content length greater than default MAX_UPLOAD_SIZE (100MB)
|
|
oversized_length = 104857600 + 1 # 100MB + 1 byte
|
|
response = client.post(
|
|
"/upload?documentId=doc-123&name=test.txt",
|
|
headers={
|
|
"Content-Length": str(oversized_length),
|
|
"Authorization": f"Bearer {OUTLINE_API_TOKEN}"
|
|
}
|
|
)
|
|
assert response.status_code == 413
|
|
assert "File too large" in response.json()["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_negative_size():
|
|
"""Verify that negative file size content lengths are rejected with 400."""
|
|
app = create_app()
|
|
client = TestClient(app)
|
|
|
|
response = client.post(
|
|
"/upload?documentId=doc-123&name=test.txt",
|
|
headers={
|
|
"Content-Length": "-10",
|
|
"Authorization": f"Bearer {OUTLINE_API_TOKEN}"
|
|
}
|
|
)
|
|
assert response.status_code == 400
|
|
assert "Invalid file size" in response.json()["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_missing_params():
|
|
"""Verify missing query params are handled with 400."""
|
|
app = create_app()
|
|
client = TestClient(app)
|
|
|
|
# Missing documentId
|
|
response = client.post(
|
|
"/upload?name=test.txt",
|
|
content=b"content",
|
|
headers={
|
|
"Content-Length": "7",
|
|
"Authorization": f"Bearer {OUTLINE_API_TOKEN}"
|
|
}
|
|
)
|
|
assert response.status_code == 400
|
|
assert "documentId" in response.json()["error"]
|
|
|
|
# Missing name
|
|
response = client.post(
|
|
"/upload?documentId=doc-123",
|
|
content=b"content",
|
|
headers={
|
|
"Content-Length": "7",
|
|
"Authorization": f"Bearer {OUTLINE_API_TOKEN}"
|
|
}
|
|
)
|
|
assert response.status_code == 400
|
|
assert "name" in response.json()["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_missing_content_length():
|
|
"""Verify that missing Content-Length header is handled correctly."""
|
|
app = create_app()
|
|
client = TestClient(app)
|
|
|
|
# We send a generator to prevent TestClient from auto-populating Content-Length
|
|
def dummy_generator():
|
|
yield b"chunk"
|
|
|
|
response = client.post(
|
|
"/upload?documentId=doc-123&name=test.txt",
|
|
content=dummy_generator(),
|
|
headers={"Authorization": f"Bearer {OUTLINE_API_TOKEN}"}
|
|
)
|
|
assert response.status_code == 400
|
|
assert "Content-Length" in response.json()["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_outline_create_failure():
|
|
"""Verify that Outline registration failures are handled correctly."""
|
|
app = create_app()
|
|
|
|
with patch.object(outline_client, "call", new_callable=AsyncMock) as mock_call:
|
|
mock_call.return_value = MOCK_ATTACH_CREATE_ERROR
|
|
|
|
client = TestClient(app)
|
|
response = client.post(
|
|
"/upload?documentId=doc-123&name=test.txt",
|
|
content=b"content",
|
|
headers={
|
|
"Content-Length": "7",
|
|
"Authorization": f"Bearer {OUTLINE_API_TOKEN}"
|
|
}
|
|
)
|
|
|
|
assert response.status_code == 502
|
|
assert "Failed to initialize attachment" in response.json()["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_endpoint_s3_failure_cleanup():
|
|
"""Verify that if the S3 upload fails, attachments.delete is called on Outline to clean up."""
|
|
app = create_app()
|
|
|
|
with patch.object(outline_client, "call", new_callable=AsyncMock) as mock_call:
|
|
mock_call.return_value = MOCK_ATTACH_CREATE_SUCCESS
|
|
|
|
# S3 PUT failure response
|
|
mock_put_response = httpx.Response(
|
|
500, request=httpx.Request("PUT", "https://s3.example.com")
|
|
)
|
|
|
|
with patch("httpx.AsyncClient") as MockClient:
|
|
mock_instance = AsyncMock()
|
|
mock_instance.put.return_value = mock_put_response
|
|
MockClient.return_value.__aenter__.return_value = mock_instance
|
|
|
|
client = TestClient(app)
|
|
response = client.post(
|
|
"/upload?documentId=doc-123&name=test.txt",
|
|
content=b"content",
|
|
headers={
|
|
"Content-Length": "7",
|
|
"Authorization": f"Bearer {OUTLINE_API_TOKEN}"
|
|
}
|
|
)
|
|
|
|
# S3 returns 500 so our gateway should return 502
|
|
assert response.status_code == 502
|
|
assert "Failed to upload attachment to storage" in response.json()["error"]
|
|
|
|
# Verify that attachments.delete was called to rollback the creation
|
|
mock_call.assert_any_call("attachments.delete", {"id": "attachment-5678-uuid"})
|