chendandan_homework/tests/test_candidate_selection_service.py

from __future__ import annotations

from datetime import date, timedelta

from app.schemas.travel_request import (
    FlightPreferences,
    HotelPreferences,
    TravelRequest,
)
from app.services.candidate_selection_service import (
    CandidateSelectionService,
)


def build_request(
    *,
    priority: str = "balanced",
    max_stops: int | None = None,
) -> TravelRequest:
    departure_date = (
        date.today() + timedelta(days=10)
    )

    return TravelRequest(
        origin_city="上海",
        destination_city="成都",
        departure_date=departure_date,
        return_date=(
            departure_date
            + timedelta(days=3)
        ),
        adults=2,
        flight_preferences=(
            FlightPreferences(
                priority=priority,
                max_stops=max_stops,
            )
        ),
        hotel_preferences=(
            HotelPreferences(
                max_price_per_night=600,
                minimum_rating=4.5,
            )
        ),
    )


def test_price_priority_prefers_lower_price() -> None:
    service = CandidateSelectionService()

    flights = [
        {
            "option_id": "cheap",
            "price": 700,
            "duration_minutes": 300,
            "stop_count": 1,
            "departure_airport": "PVG",
            "arrival_airport": "TFU",
            "departure_time": (
                "2030-01-01 08:00"
            ),
            "arrival_time": (
                "2030-01-01 13:00"
            ),
        },
        {
            "option_id": "fast",
            "price": 900,
            "duration_minutes": 120,
            "stop_count": 0,
            "departure_airport": "PVG",
            "arrival_airport": "TFU",
            "departure_time": (
                "2030-01-01 09:00"
            ),
            "arrival_time": (
                "2030-01-01 11:00"
            ),
        },
    ]

    ranked, warnings = service.rank_flights(
        flights,
        build_request(priority="price"),
    )

    assert ranked[0].option_id == "cheap"
    assert warnings == []


def test_max_stops_filters_connections() -> None:
    service = CandidateSelectionService()

    flights = [
        {
            "option_id": "direct",
            "price": 900,
            "duration_minutes": 150,
            "stop_count": 0,
        },
        {
            "option_id": "transfer",
            "price": 700,
            "duration_minutes": 280,
            "stop_count": 1,
        },
    ]

    ranked, _ = service.rank_flights(
        flights,
        build_request(max_stops=0),
    )

    assert len(ranked) == 1
    assert ranked[0].option_id == "direct"


def test_hotel_hard_filters() -> None:
    service = CandidateSelectionService()

    hotels = [
        {
            "hotel_id": "valid",
            "name": "酒店A",
            "price_per_night": 500,
            "overall_rating": 4.6,
            "location_rating": 4.5,
        },
        {
            "hotel_id": "too-expensive",
            "name": "酒店B",
            "price_per_night": 800,
            "overall_rating": 4.8,
            "location_rating": 4.7,
        },
        {
            "hotel_id": "low-rating",
            "name": "酒店C",
            "price_per_night": 400,
            "overall_rating": 4.0,
            "location_rating": 4.2,
        },
    ]

    ranked, _ = service.rank_hotels(
        hotels,
        build_request(),
    )

    assert len(ranked) == 1
    assert ranked[0].hotel_id == "valid"


def test_round_trip_uses_return_quote() -> None:
    service = CandidateSelectionService()

    combinations = [
        {
            "outbound": {
                "option_id": "out-1",
                "price": 1000,
                "duration_minutes": 180,
                "stop_count": 0,
                "currency": "CNY",
            },
            "return_flight": {
                "option_id": "ret-1",
                "price": 1800,
                "duration_minutes": 190,
                "stop_count": 0,
                "currency": "CNY",
            },
        }
    ]

    ranked = (
        service.rank_round_trip_combinations(
            combinations,
            build_request(),
        )
    )

    assert len(ranked) == 1

    # 不是1000 + 1800。
    assert ranked[0].quoted_price == 1800