test_cube_generator.py

#!/usr/bin/env python3
"""
Utility script to generate a sample .cube file and convert it to base64 for testing the LUT API.
"""

import base64
import json

def create_sample_cube_file(filename: str = "sample.cube", size: int = 8):
    """Create a simple identity LUT cube file for testing"""
    with open(filename, 'w') as f:
        f.write('TITLE "Sample Identity LUT"\n')
        f.write(f'LUT_3D_SIZE {size}\n\n')
        
        for r in range(size):
            for g in range(size):
                for b in range(size):
                    red_val = r / (size - 1)
                    green_val = g / (size - 1)
                    blue_val = b / (size - 1)
                    f.write(f'{red_val:.6f} {green_val:.6f} {blue_val:.6f}\n')
    
    print(f"Created sample cube file: {filename}")
    return filename

def cube_file_to_base64(cube_file_path: str) -> str:
    """Convert a cube file to base64 string"""
    with open(cube_file_path, 'rb') as f:
        cube_data = f.read()
    
    base64_string = base64.b64encode(cube_data).decode('utf-8')
    return base64_string

def generate_test_request(cube_file_path: str, prompt: str = "Make this LUT more cinematic with cool shadows") -> dict:
    """Generate a complete test request JSON"""
    base64_cube = cube_file_to_base64(cube_file_path)
    
    request_data = {
        "cube_file_base64": base64_cube,
        "user_prompt": prompt
    }
    
    return request_data

if __name__ == "__main__":
    cube_file = create_sample_cube_file()
    
    test_request = generate_test_request(cube_file)
    
    print(f"\nBase64 encoded cube file:")
    print(f"Length: {len(test_request['cube_file_base64'])} characters")
    print(f"First 100 chars: {test_request['cube_file_base64'][:100]}...")
    
    with open("test_request.json", "w") as f:
        json.dump(test_request, f, indent=2)
    
    print(f"\nSaved complete test request to: test_request.json")
    print(f"You can use this to test the API endpoint.")
    
    print(f"\nExample curl command:")
    print(f"curl -X POST \"http://localhost:8000/transform-lut\" \\")
    print(f"  -H \"Content-Type: application/json\" \\")
    print(f"  -d @test_request.json")