91 lines
3.3 KiB
Python
91 lines
3.3 KiB
Python
import asyncio
|
|
import os
|
|
import sys
|
|
|
|
# Add project root to path
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClient
|
|
from repos.dao import DAO
|
|
from models.Album import Album
|
|
from models.Generation import Generation, GenerationStatus
|
|
from models.enums import AspectRatios, Quality
|
|
from config import settings
|
|
|
|
# Mock config
|
|
# Use the same host as aiws.py but different DB
|
|
MONGO_HOST = settings.MONGO_HOST
|
|
DB_NAME = "bot_db_test_albums"
|
|
|
|
async def test_albums():
|
|
print(f"🚀 Starting Album Manual Verification using {MONGO_HOST}...")
|
|
|
|
# Needs to run inside a loop from main
|
|
client = AsyncIOMotorClient(MONGO_HOST)
|
|
dao = DAO(client, db_name=DB_NAME)
|
|
|
|
try:
|
|
# 1. Clean up
|
|
await client[DB_NAME]["albums"].drop()
|
|
await client[DB_NAME]["generations"].drop()
|
|
print("✅ Cleaned up test database")
|
|
|
|
# 2. Create Album
|
|
album = Album(name="Test Album", description="A test album")
|
|
print("Creating album...")
|
|
album_id = await dao.albums.create_album(album)
|
|
print(f"✅ Created Album: {album_id}")
|
|
|
|
# 3. Create Generations
|
|
gen1 = Generation(prompt="Gen 1", aspect_ratio=AspectRatios.NINESIXTEEN, quality=Quality.ONEK)
|
|
gen2 = Generation(prompt="Gen 2", aspect_ratio=AspectRatios.NINESIXTEEN, quality=Quality.ONEK)
|
|
|
|
print("Creating generations...")
|
|
gen1_id = await dao.generations.create_generation(gen1)
|
|
gen2_id = await dao.generations.create_generation(gen2)
|
|
print(f"✅ Created Generations: {gen1_id}, {gen2_id}")
|
|
|
|
# 4. Add generations to album
|
|
print("Adding generations to album...")
|
|
await dao.albums.add_generation(album_id, gen1_id)
|
|
await dao.albums.add_generation(album_id, gen2_id)
|
|
print("✅ Added generations to album")
|
|
|
|
# 5. Fetch album and check generation_ids
|
|
album_fetched = await dao.albums.get_album(album_id)
|
|
assert album_fetched is not None
|
|
assert len(album_fetched.generation_ids) == 2
|
|
assert gen1_id in album_fetched.generation_ids
|
|
assert gen2_id in album_fetched.generation_ids
|
|
print("✅ Verified generations in album")
|
|
|
|
# 6. Fetch generations by IDs via GenerationRepo
|
|
generations = await dao.generations.get_generations_by_ids([gen1_id, gen2_id])
|
|
assert len(generations) == 2
|
|
|
|
# Ensure ID type match (str vs ObjectId handling in repo)
|
|
gen_ids_fetched = [g.id for g in generations]
|
|
assert gen1_id in gen_ids_fetched
|
|
assert gen2_id in gen_ids_fetched
|
|
print("✅ Verified fetching generations by IDs")
|
|
|
|
# 7. Remove generation
|
|
print("Removing generation...")
|
|
await dao.albums.remove_generation(album_id, gen1_id)
|
|
album_fetched = await dao.albums.get_album(album_id)
|
|
assert len(album_fetched.generation_ids) == 1
|
|
assert album_fetched.generation_ids[0] == gen2_id
|
|
print("✅ Verified removing generation from album")
|
|
|
|
print("🎉 Album Verification SUCCESS")
|
|
|
|
finally:
|
|
# Cleanup client
|
|
client.close()
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(test_albums())
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|