CopilotKit/showcase/shared/python/tests/test_search_flights.py at main · codingwatching/CopilotKit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import pytest
from tools import search_flights_impl
from tools.search_flights import SURFACE_ID, CATALOG_ID

_FULL_FLIGHT = {
    "airline": "Test Air",
    "flightNumber": "TA100",
    "origin": "SFO",
    "destination": "JFK",
    "date": "Tue, Apr 15",
    "departureTime": "08:00",
    "arrivalTime": "16:00",
    "duration": "5h",
    "status": "On Time",
    "statusColor": "#22c55e",
    "price": "$299",
    "currency": "USD",
    "airlineLogo": "https://example.com/logo.png",
}


def test_returns_a2ui_operations():
    result = search_flights_impl([_FULL_FLIGHT])
    assert "a2ui_operations" in result


def test_operations_structure():
    flights = [{"airline": "Test"}]
    result = search_flights_impl(flights)
    ops = result["a2ui_operations"]
    assert any(op["type"] == "create_surface" for op in ops)
    assert any(op["type"] == "update_components" for op in ops)


def test_all_three_operation_types_present():
    result = search_flights_impl([_FULL_FLIGHT])
    ops = result["a2ui_operations"]
    types = [op["type"] for op in ops]
    assert "create_surface" in types
    assert "update_components" in types
    assert "update_data_model" in types


def test_surface_and_catalog_ids():
    result = search_flights_impl([_FULL_FLIGHT])
    ops = result["a2ui_operations"]
    create_op = next(op for op in ops if op["type"] == "create_surface")
    assert create_op["surfaceId"] == SURFACE_ID
    assert create_op["catalogId"] == CATALOG_ID


def test_flight_data_embedded_in_data_model():
    flights = [_FULL_FLIGHT, {"airline": "Second Air"}]
    result = search_flights_impl(flights)
    ops = result["a2ui_operations"]
    data_op = next(op for op in ops if op["type"] == "update_data_model")
    assert data_op["data"]["flights"] == flights


def test_empty_flights_list():
    result = search_flights_impl([])
    ops = result["a2ui_operations"]
    assert len(ops) == 3
    data_op = next(op for op in ops if op["type"] == "update_data_model")
    assert data_op["data"]["flights"] == []


def test_properly_formed_flight_objects():
    result = search_flights_impl([_FULL_FLIGHT])
    ops = result["a2ui_operations"]
    data_op = next(op for op in ops if op["type"] == "update_data_model")
    flight = data_op["data"]["flights"][0]
    for key in (
        "airline",
        "flightNumber",
        "origin",
        "destination",
        "date",
        "departureTime",
        "arrivalTime",
        "duration",
        "status",
        "price",
    ):
        assert key in flight