inpaintServer/tests/scripts/test_compatibility.py

#!/usr/bin/env python3
"""
기존 클라이언트 호환성 테스트 스크립트
기존 클라이언트 코드가 수정 없이 동작하는지 확인합니다.
"""

import os
import sys
import base64
import cv2
import numpy as np
import requests
import logging
import time
from typing import Optional

# 프로젝트 루트를 Python 경로에 추가
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))

class LegacyClientSimulator:
    """기존 클라이언트 코드를 시뮬레이션하는 클래스"""
    
    def __init__(self, base_url: str = "http://localhost:8008"):
        self.inpaint_base_url = base_url
        self.rembg_base_url = base_url
        self.inpaint_api_url = f"{base_url}/api/v1/inpaint"
        self.rembg_api_url = f"{base_url}/api/v1/remove_bg"
        
        # 로거 설정
        self.logger = logging.getLogger(__name__)
        handler = logging.StreamHandler()
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)
        self.logger.setLevel(logging.INFO)
    
    def is_server_alive(self, base_url: str) -> bool:
        """서버 상태 확인"""
        try:
            response = requests.get(f"{base_url}/health", timeout=5)
            return response.status_code == 200
        except Exception:
            return False
    
    def request_inpaint(self, image: np.ndarray, mask: np.ndarray) -> Optional[np.ndarray]:
        """기존 인페인팅 클라이언트 코드 시뮬레이션"""
        try:
            # 서버 상태 먼저 확인
            if not self.is_server_alive(self.inpaint_base_url):
                self.logger.log("인페인팅 서버가 비정상입니다. 백업 인페인팅으로 넘어갑니다.", level=logging.WARNING)
                return None

            image_data = image

            # 이미지를 base64로 인코딩
            _, img_encoded = cv2.imencode('.png', image_data)
            _, mask_encoded = cv2.imencode('.png', mask)
            img_b64 = base64.b64encode(img_encoded).decode('utf-8')
            mask_b64 = base64.b64encode(mask_encoded).decode('utf-8')
            
            # 기존 클라이언트는 model_name을 포함하지 않음
            payload = {
                "image": img_b64,
                "mask": mask_b64
            }
            
            # 기존 클라이언트는 바이너리 응답을 기대함
            response = requests.post(self.inpaint_api_url, json=payload, timeout=(5, 45))
            
            if response.status_code != 200:
                print("인페인팅 서버 에러:", response.text)
                return None
            
            # 응답이 바이너리 PNG 이미지이므로 바로 디코딩
            nparr = np.frombuffer(response.content, np.uint8)
            result = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            return result
            
        except (requests.exceptions.Timeout, requests.exceptions.ConnectionError, requests.exceptions.RequestException) as e:
            self.logger.log(f"인페인팅 서버 요청 실패: {e}, 백업 모듈로 전환", level=logging.WARNING)
            return None  # 백업으로 넘어가도록
        except Exception as e:
            self.logger.log(f"인페인팅 서버 에러: {e}", level=logging.ERROR, exc_info=True)
            return None

    def request_rembg(self, image: np.ndarray) -> Optional[np.ndarray]:
        """기존 배경 제거 클라이언트 코드 시뮬레이션"""
        try:
            # 서버 상태 먼저 확인
            if not self.is_server_alive(self.rembg_base_url):
                self.logger.log("rembg 서버가 비정상입니다. 백업 내장 rembg 모듈을 사용합니다.", level=logging.WARNING)
                return None

            # base64 인코딩
            _, img_encoded = cv2.imencode('.png', image)
            img_b64 = base64.b64encode(img_encoded).decode('utf-8')
            
            # 기존 클라이언트는 model_name을 포함하지 않음
            payload = {
                "image": img_b64
            }

            response = requests.post(self.rembg_api_url, json=payload, timeout=(5, 45))
            
            if response.status_code != 200:
                self.logger.log(f"rembg 서버 에러: {response.text}. 백업 내장 rembg 모듈을 사용합니다.", level=logging.WARNING)
                return None

            # PNG 바이너리 → numpy
            nparr = np.frombuffer(response.content, np.uint8)
            result_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)

            if result_img is None or result_img.ndim != 3:
                self.logger.log("서버 응답 이미지 디코딩 실패. 백업 내장 rembg 모듈을 사용합니다.", level=logging.WARNING)
                return None

            return result_img

        except (requests.exceptions.Timeout, requests.exceptions.ConnectionError, requests.exceptions.RequestException) as e:
            self.logger.log(f"rembg 서버 요청 실패: {e}, 백업 내장 rembg 모듈로 전환", level=logging.WARNING)
            return None
        except Exception as e:
            self.logger.log(f"request_rembg 실행 중 오류: {e}. 백업 내장 rembg 모듈을 사용합니다.", level=logging.WARNING, exc_info=True)
            return None


def create_test_images():
    """테스트용 이미지와 마스크 생성"""
    # 테스트 이미지 생성 (300x300 컬러 이미지)
    image = np.zeros((300, 300, 3), dtype=np.uint8)
    # 빨간 사각형
    cv2.rectangle(image, (50, 50), (150, 150), (0, 0, 255), -1)
    # 파란 원
    cv2.circle(image, (200, 200), 50, (255, 0, 0), -1)
    
    # 테스트 마스크 생성 (빨간 사각형 부분을 제거)
    mask = np.zeros((300, 300), dtype=np.uint8)
    cv2.rectangle(mask, (50, 50), (150, 150), 255, -1)
    
    return image, mask


def test_legacy_compatibility():
    """기존 클라이언트 호환성 테스트"""
    print("🧪 기존 클라이언트 호환성 테스트 시작")
    print("=" * 50)
    
    # 클라이언트 시뮬레이터 생성
    client = LegacyClientSimulator()
    
    # 테스트 이미지 생성
    test_image, test_mask = create_test_images()
    
    # 서버 상태 확인
    print("1. 서버 상태 확인...")
    if not client.is_server_alive("http://localhost:8008"):
        print("❌ 서버가 실행되지 않았습니다. 먼저 서버를 시작하세요.")
        print("   ./scripts/start_server.sh")
        return False
    print("✅ 서버가 정상적으로 실행 중입니다.")
    
    success_count = 0
    total_tests = 0
    
    # 인페인팅 테스트
    print("\n2. 인페인팅 API 호환성 테스트...")
    total_tests += 1
    try:
        start_time = time.time()
        result = client.request_inpaint(test_image, test_mask)
        end_time = time.time()
        
        if result is not None and isinstance(result, np.ndarray):
            print(f"✅ 인페인팅 테스트 성공 (처리 시간: {end_time - start_time:.2f}초)")
            print(f"   - 입력 이미지 크기: {test_image.shape}")
            print(f"   - 출력 이미지 크기: {result.shape}")
            print(f"   - model_name 누락 처리: ✅ (자동으로 'simple-lama' 사용)")
            print(f"   - 바이너리 PNG 응답: ✅")
            success_count += 1
        else:
            print("❌ 인페인팅 테스트 실패: 결과가 None이거나 올바르지 않습니다.")
    except Exception as e:
        print(f"❌ 인페인팅 테스트 실패: {e}")
    
    # 배경 제거 테스트
    print("\n3. 배경 제거 API 호환성 테스트...")
    total_tests += 1
    try:
        start_time = time.time()
        result = client.request_rembg(test_image)
        end_time = time.time()
        
        if result is not None and isinstance(result, np.ndarray):
            print(f"✅ 배경 제거 테스트 성공 (처리 시간: {end_time - start_time:.2f}초)")
            print(f"   - 입력 이미지 크기: {test_image.shape}")
            print(f"   - 출력 이미지 크기: {result.shape}")
            print(f"   - model_name 누락 처리: ✅ (자동으로 'rembg' 사용)")
            print(f"   - 바이너리 PNG 응답: ✅")
            success_count += 1
        else:
            print("❌ 배경 제거 테스트 실패: 결과가 None이거나 올바르지 않습니다.")
    except Exception as e:
        print(f"❌ 배경 제거 테스트 실패: {e}")
    
    # 결과 출력
    print("\n" + "=" * 50)
    print("🎯 테스트 결과")
    print("=" * 50)
    print(f"성공: {success_count}/{total_tests}")
    print(f"성공률: {success_count/total_tests*100:.1f}%")
    
    if success_count == total_tests:
        print("🎉 모든 테스트가 성공했습니다!")
        print("   기존 클라이언트 코드를 수정 없이 그대로 사용할 수 있습니다.")
        return True
    else:
        print("⚠️  일부 테스트가 실패했습니다.")
        print("   서버 로그를 확인하거나 다시 시도해보세요.")
        return False


def test_new_features():
    """새로운 기능 테스트 (선택사항)"""
    print("\n🚀 새로운 기능 테스트")
    print("=" * 50)
    
    test_image, test_mask = create_test_images()
    
    # WebP 형식 테스트
    print("1. WebP 형식 응답 테스트...")
    try:
        _, img_encoded = cv2.imencode('.png', test_image)
        _, mask_encoded = cv2.imencode('.png', test_mask)
        img_b64 = base64.b64encode(img_encoded).decode('utf-8')
        mask_b64 = base64.b64encode(mask_encoded).decode('utf-8')
        
        payload = {
            "image": img_b64,
            "mask": mask_b64,
            "model_name": "simple-lama"
        }
        
        response = requests.post(
            "http://localhost:8008/api/v1/inpaint?response_format=binary&image_format=webp", 
            json=payload, 
            timeout=(5, 45)
        )
        
        if response.status_code == 200 and response.headers.get('content-type') == 'image/webp':
            print("✅ WebP 형식 응답 성공")
        else:
            print(f"❌ WebP 형식 응답 실패: {response.status_code}")
    except Exception as e:
        print(f"❌ WebP 테스트 실패: {e}")
    
    # JSON 응답 테스트
    print("\n2. JSON 응답 형식 테스트...")
    try:
        response = requests.post(
            "http://localhost:8008/api/v1/inpaint?response_format=json", 
            json=payload, 
            timeout=(5, 45)
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get('success') and data.get('image'):
                print("✅ JSON 응답 형식 성공")
                print(f"   - 처리 시간: {data.get('processing_time', 0):.2f}초")
            else:
                print("❌ JSON 응답 형식 실패: 잘못된 응답 구조")
        else:
            print(f"❌ JSON 응답 형식 실패: {response.status_code}")
    except Exception as e:
        print(f"❌ JSON 테스트 실패: {e}")


if __name__ == "__main__":
    print("🔧 인페인팅 서버 호환성 테스트")
    print("=" * 50)
    print("이 스크립트는 기존 클라이언트 코드가 수정 없이 동작하는지 확인합니다.")
    print("")
    
    # 기존 클라이언트 호환성 테스트
    compatibility_success = test_legacy_compatibility()
    
    # 새로운 기능 테스트 (선택사항)
    if compatibility_success:
        test_new_features()
    
    print("\n" + "=" * 50)
    print("테스트 완료!")
    
    if compatibility_success:
        print("🎉 기존 클라이언트와 완벽하게 호환됩니다!")
        sys.exit(0)
    else:
        print("⚠️  호환성 문제가 발견되었습니다.")
        sys.exit(1)