test_attachment_upload.py

"""
Test script for task attachment upload API.
"""
import requests
import io
from PIL import Image

# Configuration
API_BASE_URL = "http://localhost:8003"
TASK_ID = "your-task-id-here"  # Replace with actual task ID


def get_auth_token():
    """Read JWT token from test_token.txt"""
    try:
        with open("test_token.txt", "r") as f:
            return f.read().strip()
    except FileNotFoundError:
        print("❌ test_token.txt not found. Run generate_test_token.py first.")
        exit(1)


def create_test_image(color="blue", size=(200, 200)):
    """Create a test image in memory"""
    img = Image.new('RGB', size, color=color)
    img_bytes = io.BytesIO()
    img.save(img_bytes, format='JPEG')
    img_bytes.seek(0)
    return img_bytes


def test_upload_photo(token, task_id):
    """Test uploading a photo attachment"""
    print("\n" + "="*80)
    print("TEST: Upload Photo Attachment")
    print("="*80)
    
    url = f"{API_BASE_URL}/tracker/tasks/{task_id}/attachments"
    headers = {"Authorization": f"Bearer {token}"}
    
    # Create test image
    test_image = create_test_image(color="blue")
    
    files = {
        "file": ("test_photo.jpg", test_image, "image/jpeg")
    }
    data = {
        "type": "photo"
    }
    
    print(f"Request: POST {url}")
    print(f"File: test_photo.jpg (JPEG)")
    print(f"Type: photo")
    
    response = requests.post(url, headers=headers, files=files, data=data)
    
    print(f"\nStatus Code: {response.status_code}")
    print(f"Response: {response.json()}")
    
    if response.status_code == 201:
        print("✅ Test passed: Photo uploaded successfully")
        return response.json()
    else:
        print("❌ Test failed")
        return None


def test_upload_signature(token, task_id):
    """Test uploading a signature attachment"""
    print("\n" + "="*80)
    print("TEST: Upload Signature Attachment")
    print("="*80)
    
    url = f"{API_BASE_URL}/tracker/tasks/{task_id}/attachments"
    headers = {"Authorization": f"Bearer {token}"}
    
    # Create test image
    test_image = create_test_image(color="white", size=(400, 200))
    
    files = {
        "file": ("signature.png", test_image, "image/png")
    }
    data = {
        "type": "signature"
    }
    
    print(f"Request: POST {url}")
    print(f"File: signature.png (PNG)")
    print(f"Type: signature")
    
    response = requests.post(url, headers=headers, files=files, data=data)
    
    print(f"\nStatus Code: {response.status_code}")
    print(f"Response: {response.json()}")
    
    if response.status_code == 201:
        print("✅ Test passed: Signature uploaded successfully")
        return response.json()
    else:
        print("❌ Test failed")
        return None


def test_invalid_file_type(token, task_id):
    """Test uploading invalid file type (should fail)"""
    print("\n" + "="*80)
    print("TEST: Upload Invalid File Type (should fail)")
    print("="*80)
    
    url = f"{API_BASE_URL}/tracker/tasks/{task_id}/attachments"
    headers = {"Authorization": f"Bearer {token}"}
    
    # Create a text file
    text_file = io.BytesIO(b"This is not an image")
    
    files = {
        "file": ("document.txt", text_file, "text/plain")
    }
    data = {
        "type": "photo"
    }
    
    print(f"Request: POST {url}")
    print(f"File: document.txt (text/plain)")
    print(f"Type: photo")
    
    response = requests.post(url, headers=headers, files=files, data=data)
    
    print(f"\nStatus Code: {response.status_code}")
    print(f"Response: {response.json()}")
    
    if response.status_code == 415:
        print("✅ Test passed: Invalid file type correctly rejected")
        return True
    else:
        print("❌ Test failed: Expected 415 status code")
        return False


def test_large_file(token, task_id):
    """Test uploading file that's too large (should fail)"""
    print("\n" + "="*80)
    print("TEST: Upload Large File (should fail)")
    print("="*80)
    
    url = f"{API_BASE_URL}/tracker/tasks/{task_id}/attachments"
    headers = {"Authorization": f"Bearer {token}"}
    
    # Create a large image (11MB)
    large_image = create_test_image(color="red", size=(4000, 4000))
    
    files = {
        "file": ("large_photo.jpg", large_image, "image/jpeg")
    }
    data = {
        "type": "photo"
    }
    
    print(f"Request: POST {url}")
    print(f"File: large_photo.jpg (large JPEG)")
    print(f"Type: photo")
    
    response = requests.post(url, headers=headers, files=files, data=data)
    
    print(f"\nStatus Code: {response.status_code}")
    print(f"Response: {response.json()}")
    
    if response.status_code == 413:
        print("✅ Test passed: Large file correctly rejected")
        return True
    else:
        print("⚠️  Test warning: Expected 413 status code (file might not be large enough)")
        return False


def test_invalid_task_id(token):
    """Test uploading to non-existent task (should fail)"""
    print("\n" + "="*80)
    print("TEST: Upload to Invalid Task ID (should fail)")
    print("="*80)
    
    fake_task_id = "00000000-0000-0000-0000-000000000000"
    url = f"{API_BASE_URL}/tracker/tasks/{fake_task_id}/attachments"
    headers = {"Authorization": f"Bearer {token}"}
    
    test_image = create_test_image(color="green")
    
    files = {
        "file": ("test.jpg", test_image, "image/jpeg")
    }
    data = {
        "type": "photo"
    }
    
    print(f"Request: POST {url}")
    print(f"Task ID: {fake_task_id}")
    
    response = requests.post(url, headers=headers, files=files, data=data)
    
    print(f"\nStatus Code: {response.status_code}")
    print(f"Response: {response.json()}")
    
    if response.status_code == 404:
        print("✅ Test passed: Invalid task ID correctly rejected")
        return True
    else:
        print("❌ Test failed: Expected 404 status code")
        return False


def main():
    """Run all tests"""
    print("="*80)
    print("TASK ATTACHMENT UPLOAD TEST SUITE")
    print("="*80)
    print(f"API Base URL: {API_BASE_URL}")
    print("="*80)
    
    # Get auth token
    print("\n[1/6] Getting authentication token...")
    token = get_auth_token()
    print("✅ Token loaded")
    
    # Get task ID from user
    global TASK_ID
    if TASK_ID == "your-task-id-here":
        print("\n⚠️  Please update TASK_ID in the script with a valid task ID")
        print("   You can get a task ID by running: GET /tracker/tasks/today")
        return
    
    # Test 1: Upload photo
    print("\n[2/6] Testing photo upload...")
    photo_result = test_upload_photo(token, TASK_ID)
    
    # Test 2: Upload signature
    print("\n[3/6] Testing signature upload...")
    signature_result = test_upload_signature(token, TASK_ID)
    
    # Test 3: Invalid file type
    print("\n[4/6] Testing invalid file type...")
    test_invalid_file_type(token, TASK_ID)
    
    # Test 4: Large file
    print("\n[5/6] Testing large file...")
    test_large_file(token, TASK_ID)
    
    # Test 5: Invalid task ID
    print("\n[6/6] Testing invalid task ID...")
    test_invalid_task_id(token)
    
    # Summary
    print("\n" + "="*80)
    print("TEST SUITE COMPLETED")
    print("="*80)
    
    if photo_result:
        print(f"\n✅ Photo uploaded successfully:")
        print(f"   URL: {photo_result['url']}")
        print(f"   Attachment ID: {photo_result['attachment_id']}")
    
    if signature_result:
        print(f"\n✅ Signature uploaded successfully:")
        print(f"   URL: {signature_result['url']}")
        print(f"   Attachment ID: {signature_result['attachment_id']}")
    
    print("\nTo verify in database:")
    print(f"  SELECT * FROM trans.scm_task_attachments WHERE task_id = '{TASK_ID}'::uuid;")
    print("="*80)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n❌ Test suite failed with error: {e}")
        import traceback
        traceback.print_exc()