test_app.py

import io
import pytest

from app import app, format_response, allowed_file, get_mime_type


@pytest.fixture
def client():
    app.config['TESTING'] = True
    app.config['WTF_CSRF_ENABLED'] = False
    with app.test_client() as client:
        yield client


def _get_csrf_token(client):
    """Get a CSRF token by making a GET request first."""
    response = client.get('/')
    # Extract token from the hidden input in HTML
    html = response.data.decode()
    start = html.find('name="csrf_token" value="') + len('name="csrf_token" value="')
    end = html.find('"', start)
    return html[start:end]


# --- format_response tests ---

def test_format_response_escapes_html():
    result = format_response('<script>alert("xss")</script>')
    assert '<script>' not in result
    assert '&lt;script&gt;' in result


def test_format_response_bold():
    result = format_response('**hello**')
    assert '<strong>hello</strong>' in result


def test_format_response_list_items():
    result = format_response('* item one\n* item two')
    assert '<li>item one</li>' in result
    assert '<ul>' in result


# --- allowed_file tests ---

def test_allowed_file_valid():
    assert allowed_file('photo.jpg') is True
    assert allowed_file('photo.png') is True
    assert allowed_file('photo.JPEG') is True
    assert allowed_file('photo.webp') is True


def test_allowed_file_invalid():
    assert allowed_file('script.py') is False
    assert allowed_file('doc.pdf') is False
    assert allowed_file('noext') is False


# --- get_mime_type tests ---

def test_get_mime_type_jpg():
    assert get_mime_type('photo.jpg') == 'image/jpeg'


def test_get_mime_type_png():
    assert get_mime_type('photo.png') == 'image/png'


def test_get_mime_type_unknown_fallback():
    assert get_mime_type('') == 'image/jpeg'
    assert get_mime_type('noext') == 'image/jpeg'


# --- Route tests ---

def test_get_index(client):
    response = client.get('/')
    assert response.status_code == 200
    assert b'NutriLens AI' in response.data


def test_post_without_file_redirects(client):
    token = _get_csrf_token(client)
    response = client.post('/', data={
        'user_query': 'test',
        'provider': 'ollama',
        'csrf_token': token,
    }, follow_redirects=False)
    assert response.status_code == 302


def test_post_without_csrf_returns_403(client):
    response = client.post('/', data={
        'user_query': 'test',
        'provider': 'ollama',
    })
    assert response.status_code == 403


def test_post_invalid_file_type(client):
    token = _get_csrf_token(client)
    data = {
        'user_query': 'test',
        'provider': 'ollama',
        'csrf_token': token,
        'file': (io.BytesIO(b'not an image'), 'malware.exe'),
    }
    response = client.post('/', data=data, content_type='multipart/form-data', follow_redirects=True)
    assert b'Invalid file type' in response.data


def test_post_invalid_provider(client):
    token = _get_csrf_token(client)
    data = {
        'user_query': 'test',
        'provider': 'nonexistent_provider',
        'csrf_token': token,
        'file': (io.BytesIO(b'\x89PNG\r\n'), 'photo.png'),
    }
    response = client.post('/', data=data, content_type='multipart/form-data', follow_redirects=True)
    assert b'Invalid provider' in response.data


def test_post_query_too_long(client):
    token = _get_csrf_token(client)
    data = {
        'user_query': 'x' * 501,
        'provider': 'ollama',
        'csrf_token': token,
        'file': (io.BytesIO(b'\x89PNG\r\n'), 'photo.png'),
    }
    response = client.post('/', data=data, content_type='multipart/form-data', follow_redirects=True)
    assert b'Query too long' in response.data