server.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HTTP服务端
根据请求header中的X-Case-ID字段匹配对应的响应报文并返回
"""

import json
import os
import subprocess
import sys
import time
import argparse
from flask import Flask, request, Response
from case_manager import CaseManager

app = Flask(__name__)
case_manager = CaseManager()

# 默认的case-id header字段名
CASE_ID_HEADER = 'X-Case-ID'


@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS'])
def handle_request(path):
    """处理所有HTTP请求"""
    # 获取case-id
    case_id = request.headers.get(CASE_ID_HEADER)
    
    if not case_id:
        return Response(
            json.dumps({
                'error': f'Missing required header: {CASE_ID_HEADER}',
                'message': 'Please provide a case ID in the request header'
            }),
            status=400,
            mimetype='application/json'
        )
    
    # 查找对应的用例
    case = case_manager.get_case(case_id)
    
    if not case:
        return Response(
            json.dumps({
                'error': f'Case not found: {case_id}',
                'message': f'No response case found for case ID: {case_id}'
            }),
            status=404,
            mimetype='application/json'
        )
    
    # 获取响应报文
    response_data = case.get('response', {})
    
    # 构建响应
    status_code = response_data.get('status_code', 200)
    headers = response_data.get('headers', {})
    body = response_data.get('body', '')
    
    # 处理响应体
    response_body = body
    content_type = headers.get('Content-Type', '')
    
    # 如果body是字典，转换为JSON字符串
    if isinstance(body, dict):
        response_body = json.dumps(body, ensure_ascii=False)
        if not content_type:
            content_type = 'application/json'
    elif isinstance(body, str):
        # 字符串直接使用
        if not content_type:
            content_type = 'text/plain'
    else:
        # 其他类型转换为字符串
        response_body = str(body) if body is not None else ''
        if not content_type:
            content_type = 'text/plain'
    
    # 创建响应对象
    response = Response(
        response_body,
        status=status_code,
        mimetype=content_type
    )
    
    # 设置其他响应头（排除Content-Type，因为已经通过mimetype设置）
    for key, value in headers.items():
        if key.lower() != 'content-type':
            response.headers[key] = value
    
    return response


@app.route('/health', methods=['GET'])
def health():
    """健康检查端点"""
    return {'status': 'ok', 'cases_loaded': len(case_manager.cases)}


def check_and_kill_port(port):
    """
    检查端口是否被占用，如果被占用则关闭占用该端口的进程
    
    Args:
        port: 要检查的端口号
    
    Returns:
        bool: 如果端口被占用并成功关闭返回True，如果端口未被占用返回False
    """
    try:
        # 使用lsof查找占用端口的进程
        # -ti: 只返回进程ID，不显示详细信息
        result = subprocess.run(
            ['lsof', '-ti', f':{port}'],
            capture_output=True,
            text=True
        )
        
        if result.returncode == 0 and result.stdout.strip():
            # 端口被占用，获取进程ID
            pids = result.stdout.strip().split('\n')
            print(f"Port {port} is in use by process(es): {', '.join(pids)}")
            
            # 关闭所有占用该端口的进程
            for pid in pids:
                try:
                    pid_int = int(pid.strip())
                    print(f"Killing process {pid_int}...")
                    subprocess.run(['kill', '-9', str(pid_int)], check=True)
                    print(f"Process {pid_int} killed successfully")
                except (ValueError, subprocess.CalledProcessError) as e:
                    print(f"Failed to kill process {pid}: {e}")
            
            # 等待一下，确保进程已关闭
            time.sleep(0.5)
            
            # 再次检查端口是否已释放
            result = subprocess.run(
                ['lsof', '-ti', f':{port}'],
                capture_output=True,
                text=True
            )
            
            if result.returncode == 0 and result.stdout.strip():
                print(f"Warning: Port {port} may still be in use")
                return False
            else:
                print(f"Port {port} is now available")
                return True
        else:
            # 端口未被占用
            return False
            
    except FileNotFoundError:
        # lsof命令不存在（理论上不应该发生，因为macOS/Linux都有）
        print("Warning: 'lsof' command not found, cannot check port availability")
        return False
    except Exception as e:
        print(f"Error checking port {port}: {e}")
        return False


if __name__ == '__main__':
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='HTTP服务端')
    parser.add_argument('--host', '-H', default='0.0.0.0', help='监听IP地址 (默认: 0.0.0.0)')
    parser.add_argument('--port', '-p', type=int, default=23941, help='监听端口 (默认: 32941)')
    parser.add_argument('--cases', '-c', default='cases.yaml', help='用例配置文件路径或文件夹 (默认: cases.yaml)')
    parser.add_argument('--cases-dir', help='用例文件夹路径（原始HTTP格式），可与--cases同时使用')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    
    args = parser.parse_args()
    
    HOST = args.host
    PORT = args.port
    
    # 检查并关闭占用端口的进程
    print(f"Checking port {PORT}...")
    port_was_in_use = check_and_kill_port(PORT)
    
    if port_was_in_use:
        print(f"Freed port {PORT}, starting server...")
    
    # 加载用例（支持YAML文件和文件夹）
    # 默认情况下同时加载YAML文件和cases文件夹
    if args.cases == 'cases.yaml' and not args.cases_dir:
        # 默认情况：同时加载YAML和文件夹
        case_manager.load_cases('cases.yaml')
        if os.path.exists('cases') and os.path.isdir('cases'):
            case_manager.load_cases('cases')
    else:
        # 用户指定了参数，按指定方式加载
        case_manager.load_cases(args.cases)
        if args.cases_dir:
            case_manager.load_cases(args.cases_dir)
    
    print(f"\nHTTP Server starting on http://{HOST}:{PORT}")
    print(f"Loaded {len(case_manager.cases)} cases")
    print(f"Use header '{CASE_ID_HEADER}' to specify which case to use")
    print("-" * 60)
    
    app.run(host=HOST, port=PORT, debug=args.debug)