Michal/api tests

morpheus-client/services/files.ts

@@ -16,7 +16,10 @@ export const uploadFileToServer = async (file: File, destFolder: string) => {
   }
 };
 
-export const uploadMultipleFilesToServer = async (files: File[]) => {
+export const uploadMultipleFilesToServer = async (
+  files: File[],
+  destFolder: string
+) => {
   try {
     const formData = new FormData();
     files.forEach((file) => {
@@ -29,6 +32,7 @@ export const uploadMultipleFilesToServer = async (files: File[]) => {
         headers: {
           "Content-Type": "multipart/form-data",
         },
+        params: { folder: destFolder },
       }
     );
     return { success: true, data: response.data };

morpheus-server/app/api/files_api.py

@@ -1,14 +1,12 @@
 from typing import Union
 
-from fastapi import APIRouter, Depends, File, UploadFile
-from sqlalchemy.orm import Session
-
-from morpheus_data.database.database import get_db
-
 from app.config import get_file_handlers
 from app.integrations.firebase import get_user
 from app.models.schemas import Response
 from app.services.files_services import FilesService
+from fastapi import APIRouter, Depends, File, UploadFile
+from morpheus_data.database.database import get_db
+from sqlalchemy.orm import Session
 
 router = APIRouter()
 files_repository = get_file_handlers()
@@ -17,7 +15,7 @@
 
 @router.post("/upload", response_model=Union[Response, str])
 async def upload_file_to_s3(
-    folder: str, file: UploadFile = File(...), db: Session = Depends(get_db), user=Depends(get_user)
+        folder: str, file: UploadFile = File(...), db: Session = Depends(get_db), user=Depends(get_user)
 ):
     email = user["email"]
     filename = file_service.upload_file_to_s3(file=file, db=db, email=email, folder=folder)
@@ -29,20 +27,20 @@ async def upload_file_to_s3(
 
 @router.post("/upload/multiple", response_model=Union[Response, list[str]])
 async def upload_multiple_files_to_s3(
-    files: list[UploadFile] = File(...), db: Session = Depends(get_db), user=Depends(get_user)
+        folder: str, files: list[UploadFile] = File(...), db: Session = Depends(get_db), user=Depends(get_user)
 ):
     email = user["email"]
-    filenames = file_service.upload_multiple_files_to_s3(db=db, files=files, email=email)
+    filenames = file_service.upload_multiple_files_to_s3(db=db, files=files, email=email, folder=folder)
     if not filenames:
         return {"success": False, "message": "error uploading the files"}
 
     return filenames
 
 
 @router.get("/user", response_model=Union[Response, list[str]])
-async def get_user_images(db=Depends(get_db), user=Depends(get_user)):
+async def get_user_images(folder: Union[str, None], user=Depends(get_user)):
     email = user["email"]
-    images = file_service.get_user_images(db=db, email=email)
+    images = file_service.get_user_images(email=email, folder=folder)
     if not images:
         return {"success": False, "message": "error getting the images"}
 

morpheus-server/app/api/model_category_api.py

@@ -48,7 +48,7 @@ async def update_sd_model(category: ModelCategory, db: Session = Depends(get_db)
         return Response(success=False, message=str(e))
 
 
-@router.delete("/{category_id}", response_model=Union[Response, List[ModelCategory]])
+@router.delete("/{category_id}", response_model=Union[Response, ModelCategory])
 async def delete_sd_model(category_id: UUID, db: Session = Depends(get_db)):
     try:
         category_deleted = await category_service.delete_model_category(db=db, category_id=category_id)

morpheus-server/app/api/models_api.py

@@ -60,10 +60,9 @@ async def update_sd_model(model: MLModel, db: Session = Depends(get_db)):
     return sd_model_updated
 
 
-@router.delete("/{model_source:path}", response_model=Union[Response, List[MLModel]])
+@router.delete("/{model_source:path}", response_model=Union[Response, MLModel])
 async def delete_sd_model(model_source: str, db: Session = Depends(get_db)):
     sd_model = await model_service.delete_model_by_source(db=db, model_source=model_source)
     if not sd_model:
         return Response(success=False, message="No SD Model found")
-
     return sd_model

morpheus-server/app/services/files_services.py

@@ -2,10 +2,9 @@
 from typing import Any
 
 from PIL import Image
-from sqlalchemy.orm import Session
-
 from morpheus_data.repository.files.files_interface import FileRepositoryInterface
 from morpheus_data.repository.user_repository import UserRepository
+from sqlalchemy.orm import Session
 
 logger = logging.getLogger(__name__)
 
@@ -15,12 +14,12 @@ def __init__(self, files_repository: FileRepositoryInterface):
         self.files_repository = files_repository
         self.user_repository = UserRepository()
 
-    def get_user_images(self, *, db: Session, email: str):
-        self.user_repository.get_user_data(db=db, email=email)
-        return self.files_repository.get_files(folder_name=email)
+    def get_user_images(self, *, email: str, folder: str = None):
+        destination_folder = email if folder is None else f"{folder}/{email}"
+        return self.files_repository.get_files(folder_name=destination_folder)
 
     def upload_file_to_s3(
-        self, *, db: Session, file: Any, email: str, folder: str = None, skip_validation: bool = False
+            self, *, db: Session, file: Any, email: str, folder: str = None, skip_validation: bool = False
     ):
         if not skip_validation:
             self.user_repository.get_user_data(db=db, email=email)
@@ -33,11 +32,18 @@ def upload_file_to_s3(
         else:
             logger.error("File extension not allowed")
 
-    def upload_multiple_files_to_s3(self, *, db: Session, files: list[Any], email: str):
+    # missing folder name param? (collections / avatars)
+    def upload_multiple_files_to_s3(self, *, db: Session, files: list[Any], email: str, folder: str = None):
         self.user_repository.get_user_data(db=db, email=email)
         file_urls = []
         for file in files:
-            file_url = self.upload_file_to_s3(db=db, file=file, email=email, skip_validation=True)
+            file_url = self.upload_file_to_s3(
+                db=db,
+                file=file,
+                email=email,
+                folder=folder,
+                skip_validation=True
+            )
             if file_url is not None:
                 file_urls.append(file_url)
         return file_urls

morpheus-server/requirements.txt

@@ -37,4 +37,5 @@ flower
 dynamicprompts[attentiongrabber,magicprompt]
 beautifulsoup4
 safetensors
-invisible-watermark>=0.2.0
+invisible-watermark>=0.2.0
+moto

morpheus-server/tests/conftest.py

@@ -2,57 +2,195 @@
 from httpx import AsyncClient
 
 from morpheus_data.database.database import get_db
-from morpheus_data.models.schemas import User
+from morpheus_data.models.schemas import User, CollectionCreate, ArtWorkCreate, Collection, ArtWork, Prompt, ModelCategory, MLModelCreate, MLModel, Generation
+from morpheus_data.models.models import Generation as GenerationModel
 from morpheus_data.repository.firebase_repository import FirebaseRepository
 from morpheus_data.repository.user_repository import UserRepository
+from morpheus_data.repository.collection_repository import CollectionRepository
+from morpheus_data.repository.artwork_repository import ArtWorkRepository
+from morpheus_data.repository.prompt_repository import PromptRepository
+from morpheus_data.repository.model_category_repository import ModelCategoryRepository
+from morpheus_data.repository.model_repository import ModelRepository
+from morpheus_data.repository.generation_repository import GenerationRepository
+from sqlalchemy.orm.exc import UnmappedInstanceError
+from moto import mock_s3
+import boto3
+import os
+
+from app.config import get_settings
+
+from tests.utils.prompts import generate_random_prompt
 
 from app.app import app
 
 db = next(get_db())
 
 
-class DemoUser:
+class DemoUserCredentials:
     def __init__(self):
         self.email = "[email protected]"
         self.name = "Demo User"
         self.password = "DemoPass88"
 
+@pytest.fixture(scope="session")
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+    os.environ["AWS_DEFAULT_REGION"] = "us-east-1"
+
+@pytest.fixture(scope="session")
+def mock_settings(aws_credentials):
+    return get_settings()
 
 @pytest.fixture(scope="session")
 def anyio_backend():
     return "asyncio"
 
 
 @pytest.fixture(scope="session")
-def demo_user():
-    fake_user = DemoUser()
+def demo_user_credentials():
+    fake_user = DemoUserCredentials()
     yield fake_user
 
+@pytest.fixture(scope="session")
+def demo_user(demo_user_credentials):
+    user_repository = UserRepository()
+    user = user_repository.get_user_by_email(db=db, email=demo_user_credentials.email)
+    if user is None:
+        new_user = User(email=demo_user_credentials.email, name=demo_user_credentials.name)
+        user = user_repository.create_user(db=db, user=new_user)
+    yield user
+    user_repository.delete_user(db=db, email=demo_user_credentials.email)
+
 
 @pytest.fixture(scope="session")
-async def async_app_client():
-    async with AsyncClient(app=app, base_url="https://servertest") as client:
+async def async_app_client(mock_settings):
+    def mocked_app(*args, **kwargs):
+        with mock_s3():
+            s3 = boto3.client("s3")
+            s3.create_bucket(Bucket=mock_settings.images_bucket)
+            s3.create_bucket(Bucket=mock_settings.images_temp_bucket)
+            s3.create_bucket(Bucket=mock_settings.models_bucket)
+            mocked_app = app(*args, **kwargs)
+            return mocked_app
+
+    async with AsyncClient(app=mocked_app, base_url="https://servertest") as client:
         yield client
 
-
 @pytest.fixture(scope="session")
-async def auth_header(async_app_client, demo_user):
+async def auth_header(async_app_client, demo_user_credentials, demo_user):
     firebase_repository = FirebaseRepository()
-    user = firebase_repository.get_firebase_user(email=demo_user.email)
-    if user is None:
-        firebase_repository.register_firebase_user(email=demo_user.email, password=demo_user.password)
-
-    user_repository = UserRepository()
-    user = user_repository.get_user_by_email(db=db, email=demo_user.email)
+    user = firebase_repository.get_firebase_user(email=demo_user_credentials.email)
     if user is None:
-        new_user = User(email=demo_user.email, name=demo_user.name)
-        user_repository.create_user(db=db, user=new_user)
+        firebase_repository.register_firebase_user(email=demo_user_credentials.email, password=demo_user_credentials.password)
 
     response = await firebase_repository.sign_in_with_email_and_password(
-        email=demo_user.email, password=demo_user.password
+        email=demo_user_credentials.email, password=demo_user_credentials.password
     )
     token = response["idToken"]
     yield {"Authorization": f"Bearer {token}"}
 
-    firebase_repository.remove_firebase_user(email=demo_user.email)
-    user_repository.delete_user(db=db, email=demo_user.email)
+    firebase_repository.remove_firebase_user(email=demo_user_credentials.email)
+
+@pytest.fixture(scope="function")
+def collection(demo_user):
+    collection_create = CollectionCreate(
+        name="test collection", 
+        description="test description",
+        image="https://commons.wikimedia.org/wiki/File:Trier_100_Millionen.jpg",
+    )
+    collection_repository = CollectionRepository()
+    new_collection = collection_repository.create_collection(db=db, collection=collection_create, owner=demo_user)
+    yield new_collection
+    collection_repository.delete_collection(db=db, collection_id=new_collection.id)
+
+@pytest.fixture(scope="function")
+def model_category():
+    category = ModelCategory(
+        name="test_model_category",
+        description="test_description"
+    )
+    category_repository = ModelCategoryRepository()
+    new_category = category_repository.create_category(db=db, category=category)
+    yield new_category
+    try:
+        category_repository.delete_category(db=db, category_id=new_category.id)
+    except UnmappedInstanceError:
+        pass # category deleted during test
+
+@pytest.fixture(scope="function")
+def make_artwork(demo_user, collection):
+    artworks = []
+    prompts = []
+
+    artwork_repository = ArtWorkRepository()
+    prompt_repository = PromptRepository()
+
+    def _make_artwork(
+            title: str, 
+            collection: Collection = collection,
+            image: str = "https://commons.wikimedia.org/wiki/File:Trier_100_Millionen.jpg",
+        ) -> ArtWork:
+        prompt_random = generate_random_prompt()
+        prompt = prompt_repository.get_or_create_prompt(db=db, prompt=Prompt(**prompt_random), owner=collection.owner)
+
+        artwork_create = ArtWorkCreate(
+            title = title,
+            image = image,
+            prompt = prompt,
+            collection_id = collection.id,
+        )
+        new_artwork = artwork_repository.create_artwork(db=db, artwork=artwork_create, prompt=prompt)
+        artworks.append(new_artwork)
+        prompts.append(prompt)
+        return new_artwork
+
+    yield _make_artwork
+
+    for artwork in artworks:
+        artwork_repository.delete_artwork(db=db, artwork_id=artwork.id)
+    for prompt in prompts:
+        # method not present in repository
+        #prompt_repository.delete_prompt(db=db, prompt_id=prompt.id)
+        pass
+
+@pytest.fixture(scope="function")
+def model(model_category) -> MLModel:
+    model_repository = ModelRepository()
+
+    model = MLModelCreate(**{
+            "name": "model",
+            "description": "Small dummy model",
+            "source": "hf-internal-testing/tiny-stable-diffusion-xl-pipe",
+            "kind": "diffusion",
+            "url_docs": "https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-xl-pipe",
+            "categories": [model_category],
+        }
+    )
+
+    new_model = model_repository.create_model(db=db, model=model, categories=[model_category])
+    yield new_model
+    model_repository.delete_model_by_source(db=db, model_source=new_model.source)
+
+
+@pytest.fixture(scope="function")
+def generation() -> GenerationModel:
+    generation_repository = GenerationRepository()
+    generation = generation_repository.create_generation(db=db)
+
+    generation_input = Generation(
+        id=generation.id,
+        results = [
+            "https://commons.wikimedia.org/wiki/File:Trier_100_Millionen.jpg"
+        ]
+    )
+    generation = generation_repository.update_generation(db=db, generation=generation_input)
+    return generation
+
+@pytest.fixture(scope="session")
+def test_image():
+    # load file from ./data
+    return open("tests/images/morpheus.png", "rb")