98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
import asyncio
|
|
import os
|
|
from dotenv import load_dotenv
|
|
from motor.motor_asyncio import AsyncIOMotorClient
|
|
from bson import ObjectId
|
|
|
|
# Import from project root (requires PYTHONPATH=.)
|
|
from api.service.idea_service import IdeaService
|
|
from repos.dao import DAO
|
|
from models.Idea import Idea
|
|
from models.Generation import Generation, GenerationStatus
|
|
from models.enums import AspectRatios, Quality
|
|
|
|
load_dotenv()
|
|
|
|
MONGO_HOST = os.getenv("MONGO_HOST", "mongodb://localhost:27017")
|
|
DB_NAME = os.getenv("DB_NAME", "bot_db")
|
|
|
|
print(f"Connecting to MongoDB: {MONGO_HOST}, DB: {DB_NAME}")
|
|
|
|
async def test_idea_flow():
|
|
client = AsyncIOMotorClient(MONGO_HOST)
|
|
dao = DAO(client, db_name=DB_NAME)
|
|
service = IdeaService(dao)
|
|
|
|
# 1. Create an Idea
|
|
print("Creating idea...")
|
|
user_id = "test_user_123"
|
|
project_id = "test_project_abc"
|
|
idea = await service.create_idea("My Test Idea", "Initial Description", project_id, user_id)
|
|
print(f"Idea created: {idea.id} - {idea.name}")
|
|
|
|
# 2. Update Idea
|
|
print("Updating idea...")
|
|
updated_idea = await service.update_idea(idea.id, description="Updated description")
|
|
print(f"Idea updated: {updated_idea.description}")
|
|
if updated_idea.description == "Updated description":
|
|
print("✅ Idea update successful")
|
|
else:
|
|
print("❌ Idea update FAILED")
|
|
|
|
# 3. Add Generation linked to Idea
|
|
print("Creating generation linked to idea...")
|
|
gen = Generation(
|
|
prompt="idea generation 1",
|
|
# idea_id=idea.id, <-- Intentionally NOT linking initially to test linking method
|
|
project_id=project_id,
|
|
created_by=user_id,
|
|
aspect_ratio=AspectRatios.NINESIXTEEN,
|
|
quality=Quality.ONEK,
|
|
assets_list=[]
|
|
)
|
|
gen_id = await dao.generations.create_generation(gen)
|
|
print(f"Created generation: {gen_id}")
|
|
|
|
# Link generation to idea
|
|
print("Linking generation to idea...")
|
|
success = await service.add_generation_to_idea(idea.id, gen_id)
|
|
if success:
|
|
print("✅ Linking successful")
|
|
else:
|
|
print("❌ Linking FAILED")
|
|
|
|
# Debug: Check if generation was saved with idea_id
|
|
saved_gen = await dao.generations.collection.find_one({"_id": ObjectId(gen_id)})
|
|
print(f"DEBUG: Saved Generation in DB idea_id: {saved_gen.get('idea_id')}")
|
|
|
|
# 4. Fetch Generations for Idea (Verify filtering and ordering)
|
|
print("Fetching generations for idea...")
|
|
gens = await service.dao.generations.get_generations(idea_id=idea.id) # using repo directly as service might return wrapper
|
|
print(f"Found {len(gens)} generations in idea")
|
|
|
|
if len(gens) == 1 and gens[0].id == gen_id:
|
|
print("✅ Generation retrieval successful")
|
|
else:
|
|
print("❌ Generation retrieval FAILED")
|
|
|
|
# 5. Fetch Ideas for Project
|
|
ideas = await service.get_ideas(project_id)
|
|
print(f"Found {len(ideas)} ideas for project")
|
|
|
|
# Cleaning up
|
|
print("Cleaning up...")
|
|
await service.delete_idea(idea.id)
|
|
await dao.generations.collection.delete_one({"_id": ObjectId(gen_id)})
|
|
|
|
# Verify deletion
|
|
deleted_idea = await service.get_idea(idea.id)
|
|
# IdeaRepo.delete_idea logic sets is_deleted=True
|
|
if deleted_idea and deleted_idea.is_deleted:
|
|
print(f"✅ Idea deleted successfully")
|
|
|
|
# Hard delete for cleanup
|
|
await dao.ideas.collection.delete_one({"_id": ObjectId(idea.id)})
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(test_idea_flow())
|