# python import io import os import time from config.paths import get_user_image_dir from PIL import Image as PILImage import pytest 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 token = resp.headers.get("Set-Cookie") assert token and "token=" in token 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'] = 'supersecretkey' 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