ComfyUI/tests-unit/app_test/user_manager_gallery_routes_test.py

import pytest
import os
import shutil
import json
import time
from unittest import IsolatedAsyncioTestCase # Using IsolatedAsyncioTestCase for async test methods
from unittest.mock import patch, MagicMock

from aiohttp import web

# Assuming UserManager is in app.user_manager
# Adjust the import path if your project structure is different
from app.user_manager import UserManager, GALLERY_SUFFIX

# Mock comfy.cli_args and folder_paths before they are imported by UserManager
# This is a common pattern if these modules are read at import time by the tested code.
mock_args = MagicMock()
mock_args.multi_user = False # Default to single-user mode for simplicity in most tests

mock_folder_paths = MagicMock()

# We'll set get_user_directory dynamically in test setup using tmp_path

# Apply patches at the module level if they need to be active before UserManager is imported
# or within specific test classes/methods if more fine-grained control is needed.
# For now, let's assume UserManager can be instantiated after these are patched.

@pytest.fixture
def app_client_factory(event_loop): # event_loop is a pytest-asyncio fixture
    """Factory to create aiohttp test clients."""
    async def _create_client(routes_def_func, *args_for_func):
        app = web.Application(loop=event_loop)
        routes = web.RouteTableDef()
        routes_def_func(routes, *args_for_func) # Call the function that defines routes
        app.add_routes(routes)
        return await event_loop.create_task(pytest.aiohttp.plugin.make_aiohttp_client(app))
    return _create_client


class TestUserManagerGalleryRoutes(IsolatedAsyncioTestCase):

    def setUp(self):
        # Create a temporary directory for user data
        self.test_user_dir_root = "temp_test_user_data"
        os.makedirs(self.test_user_dir_root, exist_ok=True)
        self.default_user_path = os.path.join(self.test_user_dir_root, "default")
        os.makedirs(self.default_user_path, exist_ok=True)

        # Patch folder_paths.get_user_directory and args
        self.patch_folder_paths = patch('app.user_manager.folder_paths', mock_folder_paths)
        self.patch_args = patch('app.user_manager.args', mock_args)

        self.mock_folder_paths = self.patch_folder_paths.start()
        self.mock_args = self.patch_args.start()

        self.mock_folder_paths.get_user_directory.return_value = self.test_user_dir_root
        self.mock_args.multi_user = False # Explicitly set for each test run

        self.user_manager = UserManager()

        # Setup routes for the user_manager
        self.app = web.Application()
        self.user_manager.add_routes(self.app.router)


    async def asyncSetUp(self):
        # Create a test client for making requests
        self.client = await pytest.aiohttp.plugin.make_aiohttp_client(self.app)


    async def asyncTearDown(self):
        await self.client.close() # Close the client
        self.patch_folder_paths.stop()
        self.patch_args.stop()
        if os.path.exists(self.test_user_dir_root):
            shutil.rmtree(self.test_user_dir_root)

    # --- Helper Methods ---
    def _create_file(self, filename, content="test", user="default", subdir=None):
        user_specific_path = os.path.join(self.test_user_dir_root, user)
        if subdir:
            user_specific_path = os.path.join(user_specific_path, subdir)
            os.makedirs(user_specific_path, exist_ok=True)

        filepath = os.path.join(user_specific_path, filename)
        with open(filepath, "w") as f:
            f.write(content)
        return filepath

    def _get_user_data_path(self, filename, user="default", subdir=None):
        user_specific_path = os.path.join(self.test_user_dir_root, user)
        if subdir:
            user_specific_path = os.path.join(user_specific_path, subdir)
        return os.path.join(user_specific_path, filename)

    # --- Test Cases for /gallery (GET) ---
    async def test_list_gallery_empty(self):
        resp = await self.client.get("/gallery")
        assert resp.status == 200
        data = await resp.json()
        assert data == []

    async def test_list_one_gallery_item(self):
        filename_orig = "image.png"
        filename_gallery = f"image{GALLERY_SUFFIX}.png"
        self._create_file(filename_gallery) # Create the .gallery.png file

        resp = await self.client.get("/gallery")
        assert resp.status == 200
        data = await resp.json()

        assert len(data) == 1
        item = data[0]
        assert item["filename"] == "image.png" # Original name without .gallery
        assert item["path"] == filename_gallery # Path includes .gallery
        assert item["size"] == 4 # "test"
        assert "modified" in item

    async def test_list_multiple_gallery_items_and_subdirs(self):
        self._create_file(f"img1{GALLERY_SUFFIX}.jpg")
        self._create_file(f"img2{GALLERY_SUFFIX}.jpeg", subdir="photos")
        self._create_file(f"document{GALLERY_SUFFIX}.pdf", subdir="docs/work")
        self._create_file("not_gallery.txt") # Should not be listed
        self._create_file(f"also_not_gallery{GALLERY_SUFFIX}") # No extension, should not match *.gallery.*
        self._create_file(f"another.gallery#fake.png") # Invalid char, but testing suffix rule

        resp = await self.client.get("/gallery")
        assert resp.status == 200
        data = await resp.json()

        assert len(data) == 3 # Only 3 valid gallery items
        filenames_found = sorted([item["filename"] for item in data])
        expected_filenames = sorted(["img1.jpg", "img2.jpeg", "document.pdf"])
        assert filenames_found == expected_filenames

        paths_found = sorted([item["path"] for item in data])
        expected_paths = sorted([
            f"img1{GALLERY_SUFFIX}.jpg",
            f"photos/img2{GALLERY_SUFFIX}.jpeg",
            f"docs/work/document{GALLERY_SUFFIX}.pdf"
        ])
        assert paths_found == expected_paths

    async def test_list_non_gallery_items_not_listed(self):
        self._create_file("textfile.txt")
        self._create_file(f"image_not_gallery.png") # No .gallery suffix in name
        self._create_file(f"image_with_gallery_suffix_only{GALLERY_SUFFIX}") # No further extension

        resp = await self.client.get("/gallery")
        assert resp.status == 200
        data = await resp.json()
        assert data == []

    # --- Test Cases for /userdata/{file}/gallery (POST) ---
    async def test_toggle_gallery_add(self):
        filename = "add_me.png"
        created_path = self._create_file(filename)

        resp = await self.client.post(f"/userdata/{filename}/gallery")
        assert resp.status == 200
        data = await resp.json()

        expected_new_filename = f"add_me{GALLERY_SUFFIX}.png"
        assert data["filename"] == expected_new_filename

        assert not os.path.exists(created_path)
        assert os.path.exists(self._get_user_data_path(expected_new_filename))

    async def test_toggle_gallery_add_with_subdir(self):
        filename = "add_me_subdir.jpg"
        subdir = "level1/level2"
        created_path = self._create_file(filename, subdir=subdir)

        # Path in URL needs to be URL encoded if it has slashes
        url_path = f"{subdir}/{filename}"

        resp = await self.client.post(f"/userdata/{url_path}/gallery")
        assert resp.status == 200
        data = await resp.json()

        expected_new_filename = f"add_me_subdir{GALLERY_SUFFIX}.jpg"
        assert data["filename"] == expected_new_filename # Response is basename

        assert not os.path.exists(created_path)
        assert os.path.exists(self._get_user_data_path(expected_new_filename, subdir=subdir))


    async def test_toggle_gallery_remove(self):
        original_filename_part = "remove_me"
        ext = ".jpeg"
        gallery_filename = f"{original_filename_part}{GALLERY_SUFFIX}{ext}"
        created_gallery_path = self._create_file(gallery_filename)

        resp = await self.client.post(f"/userdata/{gallery_filename}/gallery")
        assert resp.status == 200
        data = await resp.json()

        expected_new_filename = f"{original_filename_part}{ext}"
        assert data["filename"] == expected_new_filename

        assert not os.path.exists(created_gallery_path)
        assert os.path.exists(self._get_user_data_path(expected_new_filename))

    async def test_toggle_gallery_file_not_found(self):
        resp = await self.client.post("/userdata/nonexistentfile.png/gallery")
        assert resp.status == 404 # Or 400 if file not specified, but here it is specified
        data = await resp.json() # Assuming error responses are JSON
        assert "File not found" in data.get("error", "") or "File not found" in await resp.text()


    async def test_toggle_gallery_invalid_file_param(self):
        # Test with an empty file parameter or one that might be problematic
        # The route itself might catch this before user_manager logic if path is malformed
        # Depending on aiohttp's routing, this might result in a 404 for the route itself
        # or a 400 if the handler's file extraction fails.
        # UserManager's toggle_gallery_status expects `file` from `request.match_info`.
        # If `file` is empty, it returns 400 "File not specified".

        # This test is more about how aiohttp handles empty path parameters
        # For instance, a route like /userdata//gallery might not match or might pass an empty string.
        # Let's assume it passes an empty string if the route matches /userdata/{file}/gallery
        # For this, we'd need to register a route that can produce an empty 'file' match_info.
        # The current route definition /userdata/{file}/gallery will likely not match /userdata//gallery.
        # So, let's test the handler directly with a mock request if we want to ensure "File not specified".

        mock_request = MagicMock(spec=web.Request)
        mock_request.match_info = {} # No 'file'
        mock_request.headers = {} # For get_request_user_id

        # Mock get_request_user_id if it's called before file check
        # self.user_manager.get_request_user_id = MagicMock(return_value="default")

        response = await self.user_manager.toggle_gallery_status(mock_request)
        assert response.status == 400
        # text_response = await response.text() # Not needed if using response.text
        assert "File not specified" in response.text


if __name__ == "__main__":
    pytest.main()