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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
import json
from enum import Enum
from random import choice, randint
from typing import Any
from urllib.parse import urlencode
import pytest
from fastapi.encoders import jsonable_encoder
from pydantic import BaseModel, Field
from requests import Response
from starlette.testclient import TestClient
from maggma.api.API import API
from maggma.api.query_operator import NumericQuery, PaginationQuery, SparseFieldsQuery, StringQueryOperator
from maggma.api.resource import ReadOnlyResource
from maggma.stores import MemoryStore
class PetType(str, Enum):
cat = "cat"
dog = "dog"
class Owner(BaseModel):
name: str = Field(..., title="Owner's name")
age: int = Field(..., title="Owne'r Age")
weight: int = Field(..., title="Owner's weight")
class Pet(BaseModel):
name: str = Field(..., title="Pet's Name")
pet_type: PetType = Field(..., title="Pet Type")
owner_name: str = Field(..., title="Owner's name")
owners = [Owner(name=f"Person{i}", age=randint(10, 100), weight=randint(100, 200)) for i in list(range(10))]
pets = [
Pet(
name=f"Pet{i}",
pet_type=choice(list(PetType)),
owner_name=choice(owners).name,
)
for i in list(range(40))
]
@pytest.fixture()
def owner_store():
store = MemoryStore("owners", key="name")
store.connect()
store.update([jsonable_encoder(d) for d in owners])
return store
@pytest.fixture()
def pet_store():
store = MemoryStore("pets", key="name")
store.connect()
store.update([jsonable_encoder(d) for d in pets])
return store
def test_msonable(owner_store, pet_store):
owner_endpoint = ReadOnlyResource(owner_store, Owner)
pet_endpoint = ReadOnlyResource(pet_store, Pet)
manager = API({"owners": owner_endpoint, "pets": pet_endpoint})
api_dict = manager.as_dict()
for k in ["@class", "@module", "resources"]:
assert k in api_dict
def search_helper(payload, base: str = "/?", debug=True) -> tuple[Response, Any]:
"""
Helper function to directly query search endpoints
Args:
store: store f
base: base of the query, default to /query?
client: TestClient generated from FastAPI
payload: query in dictionary format
debug: True = print out the url, false don't print anything
Returns:
request.Response object that contains the response of the corresponding payload
"""
owner_store = MemoryStore("owners", key="name")
owner_store.connect()
owner_store.update([d.model_dump() for d in owners])
pets_store = MemoryStore("pets", key="name")
pets_store.connect()
pets_store.update([jsonable_encoder(d) for d in pets])
resources = {
"owners": [
ReadOnlyResource(
owner_store,
Owner,
query_operators=[
StringQueryOperator(model=Owner), # type: ignore
NumericQuery(model=Owner), # type: ignore
SparseFieldsQuery(model=Owner),
PaginationQuery(),
],
)
],
"pets": [
ReadOnlyResource(
pets_store,
Owner,
query_operators=[
StringQueryOperator(model=Pet),
NumericQuery(model=Pet),
SparseFieldsQuery(model=Pet),
PaginationQuery(),
],
)
],
}
api = API(resources=resources)
client = TestClient(api.app)
url = base + urlencode(payload)
if debug:
print(url)
res = client.get(url)
try:
data = res.json().get("data", [])
except json.decoder.JSONDecodeError:
data = res.text
return res, data
def test_cluster_run(owner_store, pet_store):
res, data = search_helper(payload="")
assert res.status_code == 200
payload = {"name": "Person1", "_limit": 10, "_all_fields": True}
res, data = search_helper(payload=payload, base="/owners/?")
assert res.status_code == 200
assert len(data) == 1
assert data[0]["name"] == "Person1"
payload = {"name": "Pet1", "_limit": 10, "_all_fields": True}
res, data = search_helper(payload=payload, base="/pets/?")
assert res.status_code == 200
assert len(data) == 1
assert data[0]["name"] == "Pet1"
|
