PenParse/penparse/webui/test/test_thumbnail_view.py

import os
from django.conf import settings
import pytest
import requests
from django.urls import reverse
from django.core.files.uploadedfile import SimpleUploadedFile
from django.urls import reverse
import requests
from ..models import ImageMemo, User


thumbnail_dir = getattr(settings, "THUMBNAIL_DIR", "thumbnails")


@pytest.fixture
def cleanup_uploaded_files():
    files_to_cleanup = []

    yield files_to_cleanup

    for image_memo, file_path in files_to_cleanup:
        if os.path.exists(file_path):
            os.remove(file_path)

        thumb_path = os.path.join(thumbnail_dir, str(image_memo.id) + ".jpg")

        if os.path.exists(thumb_path):
            os.remove(thumb_path)

        image_memo.delete()


def get_test_filepath(basename):
    return os.path.join(os.path.dirname(__file__), "data", basename)


@pytest.mark.django_db
def test_document_thumbnail_incorrect_user(client, cleanup_uploaded_files):
    user1 = User.objects.create_user(email="user1@test.com", password="password1")
    user2 = User.objects.create_user(email="user2@test.com", password="password2")

    with open(get_test_filepath("example.jpg"), "rb") as f:
        img_content = f.read()

    test_image = SimpleUploadedFile(
        name="test_image.jpg", content=img_content, content_type="image/jpeg"
    )

    image_memo = ImageMemo.objects.create(
        author=user1,
        image=test_image,
    )

    file_path = os.path.join(settings.MEDIA_ROOT, image_memo.image.name)
    cleanup_uploaded_files.append((image_memo, file_path))
    assert client.login(
        username=user2.email, password="password2"
    ), "Failed to log in test client as user2"

    url = reverse("document_thumbnail", kwargs={"pk": str(image_memo.id)})
    response = client.get(url)

    assert response.status_code == 404
    assert response.content == b"Document not found"

    assert client.login(
        username=user1.email, password="password1"
    ), "Failed to log in test client as user1"

    response = client.get(url)

    assert response.status_code == 200
    assert response["Content-Type"] == "image/jpeg"


@pytest.mark.django_db
def test_document_thumbnail_golden_path_png(client, cleanup_uploaded_files):
    user = User.objects.create_user(email="user@test.com", password="password")

    with open(get_test_filepath("example.png"), "rb") as f:
        img_content = f.read()

    test_image = SimpleUploadedFile(
        name="test_image.png", content=img_content, content_type="image/png"
    )

    image_memo = ImageMemo.objects.create(
        author=user,
        image=test_image,
    )

    file_path = os.path.join(settings.MEDIA_ROOT, image_memo.image.name)
    cleanup_uploaded_files.append((image_memo, file_path))
    assert client.login(
        username=user.email, password="password"
    ), "Failed to log in test client as user"

    url = reverse("document_thumbnail", kwargs={"pk": str(image_memo.id)})
    response = client.get(url)

    assert response.status_code == 200
    assert response["Content-Type"] == "image/jpeg"


@pytest.mark.django_db
def test_document_thumbnail_golden_path_gif(client, cleanup_uploaded_files):
    user = User.objects.create_user(email="user@test.com", password="password")

    with open(get_test_filepath("example.gif"), "rb") as f:
        img_content = f.read()

    test_image = SimpleUploadedFile(
        name="test_image.gif", content=img_content, content_type="image/gif"
    )

    image_memo = ImageMemo.objects.create(
        author=user,
        image=test_image,
    )

    file_path = os.path.join(settings.MEDIA_ROOT, image_memo.image.name)
    cleanup_uploaded_files.append((image_memo, file_path))
    assert client.login(
        username=user.email, password="password"
    ), "Failed to log in test client as user"

    url = reverse("document_thumbnail", kwargs={"pk": str(image_memo.id)})
    response = client.get(url)

    assert response.status_code == 200
    assert response["Content-Type"] == "image/jpeg"