Simple File Upload Testing

from locust import task, HttpUser
import random
import io

class FileUploadTestUser(HttpUser):
    def on_start(self):
        # File upload endpoints
        self.upload_endpoints = {
            "single": "/api/upload/single",
            "multiple": "/api/upload/multiple",
            "large": "/api/upload/large",
            "validate": "/api/upload/validate"
        }
        
        # File types for testing
        self.file_types = {
            "image": {"ext": "jpg", "mime": "image/jpeg"},
            "document": {"ext": "pdf", "mime": "application/pdf"},
            "text": {"ext": "txt", "mime": "text/plain"},
            "video": {"ext": "mp4", "mime": "video/mp4"}
        }
        
        # File sizes (in bytes)
        self.file_sizes = {
            "small": 1024,      # 1KB
            "medium": 1024*100, # 100KB
            "large": 1024*1024, # 1MB
            "xlarge": 1024*1024*10  # 10MB
        }

    def create_test_file(self, file_type="text", size="small"):
        """Create a test file in memory"""
        file_info = self.file_types[file_type]
        file_size = self.file_sizes[size]
        
        # Create file content
        if file_type == "text":
            content = b"A" * file_size
        elif file_type == "image":
            # Minimal JPEG header + data
            content = b'\xff\xd8\xff\xe0\x00\x10JFIF' + b"X" * (file_size - 14)
        else:
            content = b"TESTFILE" + b"X" * (file_size - 8)
        
        filename = f"test_file_{random.randint(1000, 9999)}.{file_info['ext']}"
        
        return {
            "content": content,
            "filename": filename,
            "mime_type": file_info["mime"]
        }

    @task(4)
    def test_single_file_upload(self):
        """Test uploading a single file"""
        file_type = random.choice(list(self.file_types.keys()))
        file_size = random.choice(["small", "medium"])
        
        test_file = self.create_test_file(file_type, file_size)
        
        files = {
            "file": (test_file["filename"], test_file["content"], test_file["mime_type"])
        }
        
        data = {
            "description": f"Test {file_type} upload",
            "category": file_type
        }
        
        with self.client.post(
            self.upload_endpoints["single"],
            files=files,
            data=data,
            name="Single File Upload"
        ) as response:
            if response.status_code == 200:
                try:
                    result = response.json()
                    file_id = result.get("file_id") or result.get("id")
                    uploaded_size = result.get("size", 0)
                    
                    print(f"Single upload success: {test_file['filename']} ({uploaded_size} bytes)")
                    
                    # Validate uploaded file size
                    if uploaded_size != len(test_file["content"]):
                        response.failure(f"Size mismatch: expected {len(test_file['content'])}, got {uploaded_size}")
                    
                except Exception as e:
                    response.failure(f"Invalid response: {str(e)}")
            elif response.status_code == 413:
                print("File too large (expected for large files)")
            elif response.status_code == 415:
                print("Unsupported file type (may be expected)")
            else:
                response.failure(f"Single file upload failed: {response.status_code}")

    @task(3)
    def test_multiple_file_upload(self):
        """Test uploading multiple files"""
        num_files = random.randint(2, 4)
        files = {}
        data = {"batch_description": "Multiple file upload test"}
        
        for i in range(num_files):
            file_type = random.choice(list(self.file_types.keys()))
            test_file = self.create_test_file(file_type, "small")
            
            files[f"file_{i}"] = (test_file["filename"], test_file["content"], test_file["mime_type"])
        
        with self.client.post(
            self.upload_endpoints["multiple"],
            files=files,
            data=data,
            name="Multiple File Upload"
        ) as response:
            if response.status_code == 200:
                try:
                    result = response.json()
                    uploaded_files = result.get("files", [])
                    
                    print(f"Multiple upload success: {len(uploaded_files)} files")
                    
                    # Validate number of uploaded files
                    if len(uploaded_files) != num_files:
                        response.failure(f"File count mismatch: expected {num_files}, got {len(uploaded_files)}")
                    
                except Exception as e:
                    response.failure(f"Invalid response: {str(e)}")
            else:
                response.failure(f"Multiple file upload failed: {response.status_code}")

    @task(2)
    def test_large_file_upload(self):
        """Test uploading large files"""
        file_size = random.choice(["large", "xlarge"])
        test_file = self.create_test_file("document", file_size)
        
        files = {
            "large_file": (test_file["filename"], test_file["content"], test_file["mime_type"])
        }
        
        data = {
            "upload_type": "large_file",
            "expected_size": len(test_file["content"])
        }
        
        with self.client.post(
            self.upload_endpoints["large"],
            files=files,
            data=data,
            name="Large File Upload",
            timeout=60  # Longer timeout for large files
        ) as response:
            if response.status_code == 200:
                try:
                    result = response.json()
                    file_id = result.get("file_id")
                    
                    print(f"Large file upload success: {test_file['filename']} ({file_size})")
                    
                except Exception as e:
                    response.failure(f"Invalid response: {str(e)}")
            elif response.status_code == 413:
                print(f"Large file rejected (size limit): {file_size}")
            elif response.status_code == 408:
                print(f"Large file upload timeout: {file_size}")
            else:
                response.failure(f"Large file upload failed: {response.status_code}")

    @task(2)
    def test_file_validation(self):
        """Test file upload validation and restrictions"""
        # Test invalid file type
        invalid_file = self.create_test_file("text", "small")
        invalid_file["filename"] = "malicious.exe"
        invalid_file["mime_type"] = "application/x-executable"
        
        files = {
            "file": (invalid_file["filename"], invalid_file["content"], invalid_file["mime_type"])
        }
        
        with self.client.post(
            self.upload_endpoints["validate"],
            files=files,
            name="Invalid File Upload"
        ) as response:
            if response.status_code == 415:
                print("Invalid file type correctly rejected")
            elif response.status_code == 400:
                print("Invalid file correctly rejected (bad request)")
            elif response.status_code == 200:
                response.failure("Invalid file type was accepted")
            else:
                print(f"Invalid file upload returned: {response.status_code}")

    @task(1)
    def test_empty_file_upload(self):
        """Test uploading empty files"""
        empty_file = {
            "content": b"",
            "filename": "empty.txt",
            "mime_type": "text/plain"
        }
        
        files = {
            "file": (empty_file["filename"], empty_file["content"], empty_file["mime_type"])
        }
        
        with self.client.post(
            self.upload_endpoints["single"],
            files=files,
            name="Empty File Upload"
        ) as response:
            if response.status_code == 400:
                print("Empty file correctly rejected")
            elif response.status_code == 200:
                print("Empty file accepted (may be valid)")
            else:
                print(f"Empty file upload returned: {response.status_code}")

    @task(1)
    def test_file_upload_limits(self):
        """Test file upload size and count limits"""
        # Try to upload many small files
        max_files = 10
        files = {}
        
        for i in range(max_files):
            test_file = self.create_test_file("text", "small")
            files[f"file_{i}"] = (test_file["filename"], test_file["content"], test_file["mime_type"])
        
        with self.client.post(
            self.upload_endpoints["multiple"],
            files=files,
            name="Upload Limit Test"
        ) as response:
            if response.status_code == 200:
                print(f"Upload limit test: {max_files} files accepted")
            elif response.status_code == 413:
                print(f"Upload limit reached (expected): {max_files} files")
            elif response.status_code == 400:
                print(f"Upload limit validation: {max_files} files rejected")
            else:
                print(f"Upload limit test returned: {response.status_code}")