chore/backend/tests/test_image_api.py

# python
import io
import os
import time
from config.paths import get_user_image_dir
from PIL import Image as PILImage
import pytest
from tests.conftest import TEST_SECRET_KEY, TEST_REFRESH_TOKEN_EXPIRY_DAYS
from werkzeug.security import generate_password_hash

from flask import Flask
from api.image_api import image_api, UPLOAD_FOLDER
from api.auth_api import auth_api
from db.db import image_db, users_db
from tinydb import Query

IMAGE_TYPE_PROFILE = 1
IMAGE_TYPE_ICON = 2
MAX_DIMENSION = 512


# Test user credentials
TEST_USER_ID = "9999999-9999-9999-9999-999999999999"
TEST_EMAIL = "testuser@example.com"
TEST_PASSWORD = "testpass"

def add_test_user():
    users_db.remove(Query().email == TEST_EMAIL)
    users_db.insert({
        "id": TEST_USER_ID,
        "first_name": "Test",
        "last_name": "User",
        "email": TEST_EMAIL,
        "password": generate_password_hash(TEST_PASSWORD),
        "verified": True,
        "image_id": "boy01"
    })

def login_and_set_cookie(client):
    resp = client.post('/auth/login', json={"email": TEST_EMAIL, "password": TEST_PASSWORD})
    assert resp.status_code == 200
    cookies = resp.headers.getlist("Set-Cookie")
    cookie_str = ' '.join(cookies)
    assert cookie_str and "access_token=" in cookie_str

def safe_remove(path):
    try:
        os.remove(path)
    except PermissionError as e:
        print(f"Warning: Could not remove {path}: {e}. Retrying...")
        time.sleep(0.1)
        try:
            os.remove(path)
        except Exception as e2:
            print(f"Warning: Still could not remove {path}: {e2}")

def remove_test_data():
    # Remove uploaded images
    user_image_dir = get_user_image_dir(TEST_USER_ID)
    if os.path.exists(user_image_dir):
        for f in os.listdir(user_image_dir):
            safe_remove(os.path.join(user_image_dir, f))
    # Clear image database
    image_db.truncate()

@pytest.fixture
def client():
    app = Flask(__name__)
    app.register_blueprint(image_api)
    app.register_blueprint(auth_api, url_prefix='/auth')
    app.config['TESTING'] = True
    app.config['SECRET_KEY'] = TEST_SECRET_KEY
    app.config['REFRESH_TOKEN_EXPIRY_DAYS'] = TEST_REFRESH_TOKEN_EXPIRY_DAYS
    with app.test_client() as c:
        add_test_user()
        remove_test_data()
        os.makedirs(get_user_image_dir(TEST_USER_ID), exist_ok=True)
        login_and_set_cookie(c)
        yield c
    for f in os.listdir(get_user_image_dir(TEST_USER_ID)):
        safe_remove(os.path.join(get_user_image_dir(TEST_USER_ID), 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()
    filename = j['filename']
    # Accept both .jpg and .jpeg extensions
    assert filename.endswith('.jpg') or filename.endswith('.jpeg'), "JPEG should be saved with .jpg or .jpeg extension"
    user_dir = get_user_image_dir(TEST_USER_ID)
    path = os.path.join(user_dir, 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()
    user_dir = get_user_image_dir(TEST_USER_ID)
    path = os.path.join(user_dir, j['filename'])
    assert os.path.exists(path)
    with PILImage.open(path) as saved:
        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()
    user_dir = get_user_image_dir(TEST_USER_ID)
    path = os.path.join(user_dir, j['filename'])
    assert os.path.exists(path)
    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), 'user_id': TEST_USER_ID}
    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', 'user_id': TEST_USER_ID}, 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', 'user_id': TEST_USER_ID}, 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', 'user_id': TEST_USER_ID}, 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', 'user_id': TEST_USER_ID}
    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