chore/tests/test_image_api.py

# python
import io
import os
from PIL import Image as PILImage
import pytest
from flask import Flask
from api.image_api import image_api, UPLOAD_FOLDER
from db.db import image_db

IMAGE_TYPE_PROFILE = 1
IMAGE_TYPE_ICON = 2
MAX_DIMENSION = 512

@pytest.fixture
def client():
    app = Flask(__name__)
    app.register_blueprint(image_api)
    app.config['TESTING'] = True
    with app.test_client() as c:
        yield c
    for f in os.listdir(UPLOAD_FOLDER):
        os.remove(os.path.join(UPLOAD_FOLDER, f))
    image_db.truncate()

def make_image_bytes(w, h, mode='RGB', color=(255, 0, 0, 255), fmt='PNG'):
    img = PILImage.new(mode, (w, h), color)
    bio = io.BytesIO()
    img.save(bio, format=fmt)
    bio.seek(0)
    return bio

def list_saved_files():
    return [f for f in os.listdir(UPLOAD_FOLDER) if os.path.isfile(os.path.join(UPLOAD_FOLDER, f))]

def test_upload_missing_type(client):
    img = make_image_bytes(100, 100, fmt='PNG')
    data = {'file': (img, 'test.png')}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 400
    assert b'Image type is required' in resp.data

def test_upload_invalid_type_value(client):
    img = make_image_bytes(50, 50, fmt='PNG')
    data = {'file': (img, 'test.png'), 'type': '3'}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 400
    assert b'Invalid image type. Must be 1 or 2' in resp.data

def test_upload_non_integer_type(client):
    img = make_image_bytes(50, 50, fmt='PNG')
    data = {'file': (img, 'test.png'), 'type': 'abc'}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 400
    assert b'Image type must be an integer' in resp.data

def test_upload_valid_png_with_id(client):
    image_db.truncate()
    img = make_image_bytes(120, 80, fmt='PNG')
    data = {'file': (img, 'sample.png'), 'type': str(IMAGE_TYPE_ICON), 'permanent': 'true'}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 200
    j = resp.get_json()
    assert 'id' in j and j['id']
    assert j['message'] == 'Image uploaded successfully'
    # DB entry
    records = image_db.all()
    assert len(records) == 1
    rec = records[0]
    assert rec['id'] == j['id']
    assert rec['permanent'] is True

def test_upload_valid_jpeg_extension_mapping(client):
    image_db.truncate()
    img = make_image_bytes(100, 60, mode='RGB', fmt='JPEG')
    data = {'file': (img, 'photo.jpg'), 'type': str(IMAGE_TYPE_PROFILE)}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 200
    j = resp.get_json()
    # Expected extension should be .jpg (this will fail with current code if format lost)
    filename = j['filename']
    assert filename.endswith('.jpg'), "JPEG should be saved with .jpg extension (code may be using None format)"
    path = os.path.join(UPLOAD_FOLDER, filename)
    assert os.path.exists(path)

def test_upload_png_alpha_preserved(client):
    image_db.truncate()
    img = make_image_bytes(64, 64, mode='RGBA', color=(10, 20, 30, 128), fmt='PNG')
    data = {'file': (img, 'alpha.png'), 'type': str(IMAGE_TYPE_ICON)}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 200
    j = resp.get_json()
    path = os.path.join(UPLOAD_FOLDER, j['filename'])
    with PILImage.open(path) as saved:
        # Alpha should exist (mode RGBA); if conversion changed it incorrectly this fails
        assert saved.mode in ('RGBA', 'LA')

def test_upload_large_image_resized(client):
    image_db.truncate()
    img = make_image_bytes(2000, 1500, fmt='PNG')
    data = {'file': (img, 'large.png'), 'type': str(IMAGE_TYPE_PROFILE)}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 200
    j = resp.get_json()
    path = os.path.join(UPLOAD_FOLDER, j['filename'])
    with PILImage.open(path) as saved:
        assert saved.width <= MAX_DIMENSION
        assert saved.height <= MAX_DIMENSION

def test_upload_invalid_image_content(client):
    bogus = io.BytesIO(b'notanimage')
    data = {'file': (bogus, 'bad.png'), 'type': str(IMAGE_TYPE_ICON)}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 400
    # Could be 'Uploaded file is not a valid image' or 'Failed to process image'
    assert b'valid image' in resp.data or b'Failed to process image' in resp.data

def test_upload_invalid_extension(client):
    image_db.truncate()
    img = make_image_bytes(40, 40, fmt='PNG')
    data = {'file': (img, 'note.gif'), 'type': str(IMAGE_TYPE_ICON)}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 400
    assert b'Invalid file type' in resp.data or b'Invalid file type or no file selected' in resp.data

def test_request_image_success(client):
    image_db.truncate()
    img = make_image_bytes(30, 30, fmt='PNG')
    data = {'file': (img, 'r.png'), 'type': str(IMAGE_TYPE_ICON)}
    up = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert up.status_code == 200
    recs = image_db.all()
    image_id = recs[0]['id']
    resp = client.get(f'/image/request/{image_id}')
    assert resp.status_code == 200

def test_request_image_not_found(client):
    resp = client.get('/image/request/missing-id')
    assert resp.status_code == 404
    assert b'Image not found' in resp.data

def test_list_images_filter_type(client):
    image_db.truncate()
    # Upload type 1
    for _ in range(2):
        img = make_image_bytes(20, 20, fmt='PNG')
        client.post('/image/upload', data={'file': (img, 'a.png'), 'type': '1'}, content_type='multipart/form-data')
    # Upload type 2
    for _ in range(3):
        img = make_image_bytes(25, 25, fmt='PNG')
        client.post('/image/upload', data={'file': (img, 'b.png'), 'type': '2'}, content_type='multipart/form-data')

    resp = client.get('/image/list?type=2')
    assert resp.status_code == 200
    j = resp.get_json()
    assert j['count'] == 3

def test_list_images_invalid_type_query(client):
    resp = client.get('/image/list?type=99')
    assert resp.status_code == 400
    assert b'Invalid image type' in resp.data

def test_list_images_all(client):
    image_db.truncate()
    for _ in range(4):
        img = make_image_bytes(10, 10, fmt='PNG')
        client.post('/image/upload', data={'file': (img, 'x.png'), 'type': '2'}, content_type='multipart/form-data')
    resp = client.get('/image/list')
    assert resp.status_code == 200
    j = resp.get_json()
    assert j['count'] == 4
    assert len(j['ids']) == 4

def test_permanent_flag_false_default(client):
    image_db.truncate()
    img = make_image_bytes(32, 32, fmt='PNG')
    data = {'file': (img, 't.png'), 'type': '1'}
    resp = client.post('/image/upload', data=data, content_type='multipart/form-data')
    assert resp.status_code == 200
    recs = image_db.all()
    assert recs[0]['permanent'] is False