Skip to content

Commit 6d7fca0

Browse files
more esql DSL tests
1 parent 21a1572 commit 6d7fca0

File tree

2 files changed

+146
-6
lines changed

2 files changed

+146
-6
lines changed

test_elasticsearch/test_dsl/test_integration/_async/test_esql.py

Lines changed: 73 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -145,9 +145,9 @@ async def test_esql(async_client):
145145
# get the full names of the employees
146146
query = (
147147
ESQL.from_(Employee)
148-
.eval(name=functions.concat(Employee.first_name, " ", Employee.last_name))
149-
.keep("name")
150-
.sort("name")
148+
.eval(full_name=functions.concat(Employee.first_name, " ", Employee.last_name))
149+
.keep("full_name")
150+
.sort("full_name")
151151
.limit(10)
152152
)
153153
r = await async_client.esql.query(query=str(query))
@@ -182,3 +182,73 @@ async def test_esql(async_client):
182182
)
183183
r = await async_client.esql.query(query=str(query), params=["Maria"])
184184
assert r.body["values"] == [["Luna"], ["Cannon"]]
185+
186+
187+
@pytest.mark.asyncio
188+
async def test_esql_dsl(async_client):
189+
await load_db()
190+
191+
# get employees with first name "Maria"
192+
query = (
193+
Employee.esql_from()
194+
.where(Employee.first_name == "Maria")
195+
.sort("last_name")
196+
.limit(10)
197+
)
198+
marias = []
199+
async for emp in Employee.esql_execute(query):
200+
marias.append(emp)
201+
assert len(marias) == 2
202+
assert marias[0].last_name == "Cannon"
203+
assert marias[0].address.address == "322 NW Johnston"
204+
assert marias[0].address.city == "Bakerburgh, MP"
205+
assert marias[1].last_name == "Luna"
206+
assert marias[1].address.address == "5861 Morgan Springs"
207+
assert marias[1].address.city == "Lake Daniel, WI"
208+
209+
# run a query with a missing field
210+
query = (
211+
Employee.esql_from()
212+
.where(Employee.first_name == "Maria")
213+
.drop(Employee.address.city)
214+
.sort("last_name")
215+
.limit(10)
216+
)
217+
with pytest.raises(ValueError):
218+
await Employee.esql_execute(query).__anext__()
219+
marias = []
220+
async for emp in Employee.esql_execute(query, ignore_missing_fields=True):
221+
marias.append(emp)
222+
assert marias[0].last_name == "Cannon"
223+
assert marias[0].address.address == "322 NW Johnston"
224+
assert marias[0].address.city is None
225+
assert marias[1].last_name == "Luna"
226+
assert marias[1].address.address == "5861 Morgan Springs"
227+
assert marias[1].address.city is None
228+
229+
# run a query with additional calculated fields
230+
query = (
231+
Employee.esql_from()
232+
.where(Employee.first_name == "Maria")
233+
.eval(
234+
full_name=functions.concat(Employee.first_name, " ", Employee.last_name),
235+
height_cm=functions.to_integer(Employee.height * 100),
236+
)
237+
.sort("last_name")
238+
.limit(10)
239+
)
240+
assert isinstance(await Employee.esql_execute(query).__anext__(), Employee)
241+
assert isinstance(
242+
await Employee.esql_execute(query, return_additional=True).__anext__(), tuple
243+
)
244+
marias = []
245+
async for emp, extra in Employee.esql_execute(query, return_additional=True):
246+
marias.append([emp, extra])
247+
assert marias[0][0].last_name == "Cannon"
248+
assert marias[0][0].address.address == "322 NW Johnston"
249+
assert marias[0][0].address.city == "Bakerburgh, MP"
250+
assert marias[0][1] == {"full_name": "Maria Cannon", "height_cm": 208}
251+
assert marias[1][0].last_name == "Luna"
252+
assert marias[1][0].address.address == "5861 Morgan Springs"
253+
assert marias[1][0].address.city == "Lake Daniel, WI"
254+
assert marias[1][1] == {"full_name": "Maria Luna", "height_cm": 189}

test_elasticsearch/test_dsl/test_integration/_sync/test_esql.py

Lines changed: 73 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -145,9 +145,9 @@ def test_esql(client):
145145
# get the full names of the employees
146146
query = (
147147
ESQL.from_(Employee)
148-
.eval(name=functions.concat(Employee.first_name, " ", Employee.last_name))
149-
.keep("name")
150-
.sort("name")
148+
.eval(full_name=functions.concat(Employee.first_name, " ", Employee.last_name))
149+
.keep("full_name")
150+
.sort("full_name")
151151
.limit(10)
152152
)
153153
r = client.esql.query(query=str(query))
@@ -182,3 +182,73 @@ def test_esql(client):
182182
)
183183
r = client.esql.query(query=str(query), params=["Maria"])
184184
assert r.body["values"] == [["Luna"], ["Cannon"]]
185+
186+
187+
@pytest.mark.sync
188+
def test_esql_dsl(client):
189+
load_db()
190+
191+
# get employees with first name "Maria"
192+
query = (
193+
Employee.esql_from()
194+
.where(Employee.first_name == "Maria")
195+
.sort("last_name")
196+
.limit(10)
197+
)
198+
marias = []
199+
for emp in Employee.esql_execute(query):
200+
marias.append(emp)
201+
assert len(marias) == 2
202+
assert marias[0].last_name == "Cannon"
203+
assert marias[0].address.address == "322 NW Johnston"
204+
assert marias[0].address.city == "Bakerburgh, MP"
205+
assert marias[1].last_name == "Luna"
206+
assert marias[1].address.address == "5861 Morgan Springs"
207+
assert marias[1].address.city == "Lake Daniel, WI"
208+
209+
# run a query with a missing field
210+
query = (
211+
Employee.esql_from()
212+
.where(Employee.first_name == "Maria")
213+
.drop(Employee.address.city)
214+
.sort("last_name")
215+
.limit(10)
216+
)
217+
with pytest.raises(ValueError):
218+
Employee.esql_execute(query).__next__()
219+
marias = []
220+
for emp in Employee.esql_execute(query, ignore_missing_fields=True):
221+
marias.append(emp)
222+
assert marias[0].last_name == "Cannon"
223+
assert marias[0].address.address == "322 NW Johnston"
224+
assert marias[0].address.city is None
225+
assert marias[1].last_name == "Luna"
226+
assert marias[1].address.address == "5861 Morgan Springs"
227+
assert marias[1].address.city is None
228+
229+
# run a query with additional calculated fields
230+
query = (
231+
Employee.esql_from()
232+
.where(Employee.first_name == "Maria")
233+
.eval(
234+
full_name=functions.concat(Employee.first_name, " ", Employee.last_name),
235+
height_cm=functions.to_integer(Employee.height * 100),
236+
)
237+
.sort("last_name")
238+
.limit(10)
239+
)
240+
assert isinstance(Employee.esql_execute(query).__next__(), Employee)
241+
assert isinstance(
242+
Employee.esql_execute(query, return_additional=True).__next__(), tuple
243+
)
244+
marias = []
245+
for emp, extra in Employee.esql_execute(query, return_additional=True):
246+
marias.append([emp, extra])
247+
assert marias[0][0].last_name == "Cannon"
248+
assert marias[0][0].address.address == "322 NW Johnston"
249+
assert marias[0][0].address.city == "Bakerburgh, MP"
250+
assert marias[0][1] == {"full_name": "Maria Cannon", "height_cm": 208}
251+
assert marias[1][0].last_name == "Luna"
252+
assert marias[1][0].address.address == "5861 Morgan Springs"
253+
assert marias[1][0].address.city == "Lake Daniel, WI"
254+
assert marias[1][1] == {"full_name": "Maria Luna", "height_cm": 189}

0 commit comments

Comments
 (0)