chendandan_homework/scripts/test_travel_service.py

from __future__ import annotations

"""测试TravelSearchService的航班/酒店搜索与规范化功能。"""

import asyncio
from collections import Counter
from datetime import date

from rich.console import Console
from rich.table import Table

from app.clients.serpapi_client import (
    SerpApiClient,
    SerpApiError,
)
from app.config import get_settings
from app.schemas.flight import FlightSearchQuery
from app.schemas.hotel import HotelSearchQuery
from app.services.travel_search_service import (
    TravelSearchService,
)


console = Console()


def format_duration(minutes: int | None) -> str:
    if minutes is None:
        return "未知"

    hours, remaining_minutes = divmod(minutes, 60)

    return f"{hours}小时{remaining_minutes}分"


def format_flight_numbers(flight) -> str:
    flight_numbers = [
        segment.flight_number
        for segment in flight.segments
        if segment.flight_number
    ]

    return " / ".join(flight_numbers) or "未知"


async def test_flight_service(
    service: TravelSearchService,
) -> None:
    console.rule("[bold]TravelSearchService 航班测试")

    query = FlightSearchQuery(
        departure_airports=["SHA", "PVG"],
        arrival_airports=["CTU", "TFU"],
        outbound_date=date(2026, 8, 10),
        flight_type="one_way",
        adults=2,

        # 当前不限制时间、机场、中转次数和价格。
        show_hidden=False,
        deep_search=False,
    )

    result = await service.search_flights(query)

    airport_counts = Counter(
        flight.final_arrival_airport_code
        for flight in result.flights
    )

    console.print(
        {
            "标准化航班总数": result.total_count,
            "抵达机场分布": dict(airport_counts),
            "数据源": result.source.provider,
            "查询ID": result.source.search_id,
        }
    )

    table = Table(
        title=(
            f"航班结果预览：展示10条，"
            f"实际保留{result.total_count}条"
        )
    )

    table.add_column("来源")
    table.add_column("航司/航班号")
    table.add_column("出发")
    table.add_column("到达")
    table.add_column("耗时")
    table.add_column("中转")
    table.add_column("单张票价")

    # 这里只限制终端预览，不修改 result.flights。
    for flight in result.flights[:10]:
        first_segment = flight.segments[0]
        last_segment = flight.segments[-1]

        airline_text = (
            "/".join(flight.airlines)
            if flight.airlines
            else "未知"
        )

        table.add_row(
            flight.source_group,
            (
                f"{airline_text}\n"
                f"{format_flight_numbers(flight)}"
            ),
            (
                f"{first_segment.departure_airport_code}\n"
                f"{flight.departure_time:%m-%d %H:%M}"
            ),
            (
                f"{last_segment.arrival_airport_code}\n"
                f"{flight.arrival_time:%m-%d %H:%M}"
            ),
            format_duration(
                flight.total_duration_minutes
            ),
            (
                "直飞"
                if flight.stop_count == 0
                else f"{flight.stop_count}次"
            ),
            (
                f"¥{flight.price}"
                if flight.price is not None
                else "未知"
            ),
        )

    console.print(table)

    for warning in result.warnings:
        console.print(f"[yellow]提示：{warning}[/yellow]")


async def test_hotel_service(
    service: TravelSearchService,
) -> None:
    console.rule("[bold]TravelSearchService 酒店测试")

    query = HotelSearchQuery(
        query="成都酒店",
        check_in_date=date(2026, 8, 10),
        check_out_date=date(2026, 8, 13),
        adults=2,
    )

    result = await service.search_hotels(query)

    console.print(
        {
            "标准化酒店总数": result.total_count,
            "数据源": result.source.provider,
            "查询ID": result.source.search_id,
        }
    )

    table = Table(
        title=(
            f"酒店结果预览：展示10家，"
            f"实际保留{result.total_count}家"
        )
    )

    table.add_column("酒店")
    table.add_column("星级")
    table.add_column("评分")
    table.add_column("每晚")
    table.add_column("整个入住期")
    table.add_column("客观信息")

    # 这里是客观信息摘要，不是个性化推荐理由。
    for hotel in result.hotels[:10]:
        highlights: list[str] = []

        if hotel.review_count is not None:
            highlights.append(
                f"{hotel.review_count}条评价"
            )

        if hotel.location_rating is not None:
            highlights.append(
                f"位置评分{hotel.location_rating}"
            )

        highlights.extend(hotel.amenities[:2])

        if hotel.deal_description:
            highlights.append(
                hotel.deal_description
            )

        table.add_row(
            hotel.name,
            (
                f"{hotel.hotel_class}星"
                if hotel.hotel_class
                else "未知"
            ),
            (
                str(hotel.overall_rating)
                if hotel.overall_rating is not None
                else "未知"
            ),
            (
                f"¥{hotel.price_per_night:.0f}"
                if hotel.price_per_night is not None
                else "未知"
            ),
            (
                f"¥{hotel.total_price:.0f}"
                if hotel.total_price is not None
                else "未知"
            ),
            "；".join(highlights) or "暂无",
        )

    console.print(table)


async def main() -> None:
    settings = get_settings()

    async with SerpApiClient(
        api_key=settings.require(
            "serpapi_api_key"
        ),
        timeout_seconds=(
            settings.request_timeout_seconds
        ),
    ) as client:
        service = TravelSearchService(client)

        await test_flight_service(service)
        await test_hotel_service(service)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except SerpApiError as exc:
        console.print(
            f"[red]真实接口测试失败：{exc}[/red]"
        )
        raise SystemExit(1) from exc