site11/services/pipeline/test_pipeline.py

#!/usr/bin/env python3
"""
Pipeline Test Script
파이프라인 전체 플로우를 테스트하는 스크립트
"""
import asyncio
import json
from datetime import datetime
from motor.motor_asyncio import AsyncIOMotorClient
import redis.asyncio as redis
from shared.models import KeywordSubscription, PipelineJob

async def test_pipeline():
    """파이프라인 테스트"""
    
    # MongoDB 연결
    mongo_client = AsyncIOMotorClient("mongodb://mongodb:27017")
    db = mongo_client.pipeline
    
    # Redis 연결
    redis_client = redis.Redis(host='redis', port=6379, decode_responses=True)
    
    # 1. 테스트 키워드 추가
    test_keyword = KeywordSubscription(
        keyword="전기차",
        language="ko",
        schedule="*/1 * * * *",  # 1분마다 (테스트용)
        is_active=True,
        is_priority=True,
        rss_feeds=[
            "https://news.google.com/rss/search?q=전기차&hl=ko&gl=KR&ceid=KR:ko",
            "https://news.google.com/rss/search?q=electric+vehicle&hl=en&gl=US&ceid=US:en"
        ],
        categories=["technology", "automotive", "environment"],
        owner="test_user"
    )
    
    # MongoDB에 저장
    await db.keyword_subscriptions.replace_one(
        {"keyword": test_keyword.keyword},
        test_keyword.dict(),
        upsert=True
    )
    print(f"✅ 키워드 '{test_keyword.keyword}' 추가 완료")
    
    # 2. 즉시 파이프라인 트리거 (스케줄러를 거치지 않고 직접)
    job = PipelineJob(
        keyword_id=test_keyword.keyword_id,
        keyword=test_keyword.keyword,
        stage="rss_collection",
        data={
            "rss_feeds": test_keyword.rss_feeds,
            "categories": test_keyword.categories
        },
        priority=1 if test_keyword.is_priority else 0
    )
    
    # Redis 큐에 직접 추가 (QueueMessage 형식으로)
    from shared.queue_manager import QueueMessage
    message = QueueMessage(
        queue_name="rss_collection",
        job=job
    )
    await redis_client.lpush("queue:rss_collection", message.json())
    print(f"✅ 작업을 RSS Collection 큐에 추가: {job.job_id}")
    
    # 3. 파이프라인 상태 모니터링
    print("\n📊 파이프라인 실행 모니터링 중...")
    print("각 단계별 로그를 확인하려면 다음 명령을 실행하세요:")
    print("  docker-compose logs -f pipeline-rss-collector")
    print("  docker-compose logs -f pipeline-google-search")
    print("  docker-compose logs -f pipeline-ai-summarizer")
    print("  docker-compose logs -f pipeline-translator")
    print("  docker-compose logs -f pipeline-image-generator")
    print("  docker-compose logs -f pipeline-article-assembly")
    
    # 큐 상태 확인
    for i in range(10):
        await asyncio.sleep(5)
        
        # 각 큐의 길이 확인
        queues = [
            "queue:rss_collection",
            "queue:google_search", 
            "queue:ai_summarization",
            "queue:translation",
            "queue:image_generation",
            "queue:article_assembly"
        ]
        
        print(f"\n[{datetime.now().strftime('%H:%M:%S')}] 큐 상태:")
        for queue in queues:
            length = await redis_client.llen(queue)
            if length > 0:
                print(f"  {queue}: {length} 작업 대기 중")
    
    # 4. 최종 결과 확인
    print("\n📄 MongoDB에서 생성된 기사 확인 중...")
    articles = await db.articles.find({"keyword": test_keyword.keyword}).to_list(length=5)
    
    if articles:
        print(f"✅ {len(articles)}개의 기사 생성 완료!")
        for article in articles:
            print(f"\n제목: {article.get('title', 'N/A')}")
            print(f"ID: {article.get('article_id', 'N/A')}")
            print(f"생성 시간: {article.get('created_at', 'N/A')}")
            print(f"처리 시간: {article.get('processing_time', 'N/A')}초")
            print(f"이미지 수: {len(article.get('images', []))}")
    else:
        print("⚠️ 아직 기사가 생성되지 않았습니다. 조금 더 기다려주세요.")
    
    # 연결 종료
    await redis_client.close()
    mongo_client.close()

if __name__ == "__main__":
    print("🚀 파이프라인 테스트 시작")
    asyncio.run(test_pipeline())