chendandan_homework/scripts/test_serpapi_client.py

from __future__ import annotations
"""测试SerpApiClient封装层的航班/酒店查询与数据规范化。"""
import asyncio
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


console = Console()


def _format_duration(minutes: object) -> str:
    """把分钟数转成可读时长，例如 200 → 3小时20分。"""
    if not isinstance(minutes, (int, float)):
        return "?"
    total = int(minutes)
    hours = total // 60
    mins = total % 60
    if hours == 0:
        return f"{mins}分钟"
    if mins == 0:
        return f"{hours}小时"
    return f"{hours}小时{mins}分"


def _format_stops(segments_count: int) -> str:
    """根据航段数返回直飞/中转描述。"""
    if segments_count <= 1:
        return "直飞"
    return f"中转{segments_count - 1}次"


def _build_flight_route(segments: list[dict]) -> str:
    """从航段列表拼接路线，例如 PVG → CTU。"""
    if not segments:
        return "未知路线"

    codes: list[str] = []
    for i, seg in enumerate(segments):
        dep = seg.get("departure_airport", {}).get("id", "?")
        arr = seg.get("arrival_airport", {}).get("id", "?")
        if i == 0:
            codes.append(dep)
        codes.append(arr)
    return " → ".join(codes)


async def test_flights(
    client: SerpApiClient,
) -> None:
    console.rule("[bold]测试SerpApiClient航班查询")

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

    data = await client.search_flights_raw(query)

    best_flights = data.get("best_flights", [])
    other_flights = data.get("other_flights", [])

    all_flights = [
        *best_flights,
        *other_flights,
    ]

    console.print(
        {
            "查询人数(adults)": query.adults,
            "best_flights数量": len(best_flights),
            "other_flights数量": len(other_flights),
            "全部航班数量": len(all_flights),
            "说明": "price是SerpApi返回的总价，人均=总价/adults",
        }
    )

    table = Table(title="航班查询结果预览")
    table.add_column("#", justify="right", style="dim")
    table.add_column("来源", style="bold")
    table.add_column("航空公司 / 航班号")
    table.add_column("出发机场 · 起飞时间")
    table.add_column("到达机场 · 落地时间")
    table.add_column("总时长")
    table.add_column("直飞/中转")
    table.add_column("总价")
    table.add_column("人均价格")

    seq = 0
    for source_group, options in (
        ("best_flights", best_flights),
        ("other_flights", other_flights),
    ):
        for option in options[:3]:
            segments = option.get("flights", [])

            if not segments:
                continue

            seq += 1

            # --- 航空公司 + 航班号 ---
            airlines_flight_numbers: list[str] = []
            for seg in segments:
                airline = seg.get("airline", "?")
                fn_ = seg.get("flight_number", "?")
                airlines_flight_numbers.append(f"{airline} {fn_}")
            airline_flight_str = "\n".join(airlines_flight_numbers)

            # --- 出发 ---
            first_seg = segments[0]
            dep_airport = first_seg.get("departure_airport", {})
            dep_code = dep_airport.get("id", "?")
            dep_time = dep_airport.get("time", "?")
            departure_str = f"{dep_code}\n{dep_time}"

            # --- 到达 ---
            last_seg = segments[-1]
            arr_airport = last_seg.get("arrival_airport", {})
            arr_code = arr_airport.get("id", "?")
            arr_time = arr_airport.get("time", "?")
            arrival_str = f"{arr_code}\n{arr_time}"

            # --- 总时长 ---
            total_dur = option.get("total_duration", "?")

            # --- 直飞/中转 ---
            stops_str = _format_stops(len(segments))

            # --- 总价 & 人均 ---
            price_raw = option.get("price")
            if isinstance(price_raw, (int, float)):
                total_price = int(price_raw)
                per_person = round(total_price / query.adults)
                price_str = f"¥{total_price}"
                per_person_str = f"¥{per_person}"
            else:
                price_str = "?"
                per_person_str = "?"

            # --- 路线概要 ---
            route = _build_flight_route(segments)

            table.add_row(
                str(seq),
                f"{source_group}\n{route}",
                airline_flight_str,
                departure_str,
                arrival_str,
                _format_duration(total_dur),
                stops_str,
                price_str,
                per_person_str,
            )

    console.print(table)


async def test_hotels(
    client: SerpApiClient,
) -> None:
    console.rule("[bold]测试SerpApiClient酒店查询")

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

    data = await client.search_hotels_raw(query)

    hotels = data.get("properties", [])

    console.print(
        {
            "全部酒店数量": len(hotels),
        }
    )

    table = Table(title="酒店查询结果预览")
    table.add_column("名称")
    table.add_column("评分")
    table.add_column("每晚价格")
    table.add_column("星级")

    # 这里只限制终端显示数量。
    for hotel in hotels[:5]:
        rate = hotel.get("rate_per_night", {})

        hotel_class = (
            hotel.get("extracted_hotel_class")
            or hotel.get("hotel_class")
        )

        table.add_row(
            str(hotel.get("name", "?")),
            str(hotel.get("overall_rating", "?")),
            str(rate.get("extracted_lowest", "?")),
            str(hotel_class or "?"),
        )

    console.print(table)


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

    api_key = settings.require(
        "serpapi_api_key"
    )

    async with SerpApiClient(
        api_key=api_key,
        timeout_seconds=(
            settings.request_timeout_seconds
        ),
    ) as client:
        await test_flights(client)
        await test_hotels(client)


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