Skip to content

Commit 728ab41

Browse files
nisha2003nisha2003
committed
ruff
1 parent 426398b commit 728ab41

1 file changed

Lines changed: 92 additions & 15 deletions

File tree

‎tests/databricks_ai_bridge/test_genie.py‎

Lines changed: 92 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -211,8 +211,20 @@ def test_parse_query_result_with_data():
211211
},
212212
"result": {
213213
"data_array": [
214-
{"values": [{"string_value": "1"}, {"string_value": "Alice"}, {"string_value": "2023-10-01T00:00:00Z"}]},
215-
{"values": [{"string_value": "2"}, {"string_value": "Bob"}, {"string_value": "2023-10-02T00:00:00Z"}]},
214+
{
215+
"values": [
216+
{"string_value": "1"},
217+
{"string_value": "Alice"},
218+
{"string_value": "2023-10-01T00:00:00Z"},
219+
]
220+
},
221+
{
222+
"values": [
223+
{"string_value": "2"},
224+
{"string_value": "Bob"},
225+
{"string_value": "2023-10-02T00:00:00Z"},
226+
]
227+
},
216228
]
217229
},
218230
}
@@ -272,16 +284,76 @@ def test_parse_query_result_trims_data(truncate_results):
272284
},
273285
"result": {
274286
"data_array": [
275-
{"values": [{"string_value": "1"}, {"string_value": "Alice"}, {"string_value": "2023-10-01T00:00:00Z"}]},
276-
{"values": [{"string_value": "2"}, {"string_value": "Bob"}, {"string_value": "2023-10-02T00:00:00Z"}]},
277-
{"values": [{"string_value": "3"}, {"string_value": "Charlie"}, {"string_value": "2023-10-03T00:00:00Z"}]},
278-
{"values": [{"string_value": "4"}, {"string_value": "David"}, {"string_value": "2023-10-04T00:00:00Z"}]},
279-
{"values": [{"string_value": "5"}, {"string_value": "Eve"}, {"string_value": "2023-10-05T00:00:00Z"}]},
280-
{"values": [{"string_value": "6"}, {"string_value": "Frank"}, {"string_value": "2023-10-06T00:00:00Z"}]},
281-
{"values": [{"string_value": "7"}, {"string_value": "Grace"}, {"string_value": "2023-10-07T00:00:00Z"}]},
282-
{"values": [{"string_value": "8"}, {"string_value": "Hank"}, {"string_value": "2023-10-08T00:00:00Z"}]},
283-
{"values": [{"string_value": "9"}, {"string_value": "Ivy"}, {"string_value": "2023-10-09T00:00:00Z"}]},
284-
{"values": [{"string_value": "10"}, {"string_value": "Jack"}, {"string_value": "2023-10-10T00:00:00Z"}]},
287+
{
288+
"values": [
289+
{"string_value": "1"},
290+
{"string_value": "Alice"},
291+
{"string_value": "2023-10-01T00:00:00Z"},
292+
]
293+
},
294+
{
295+
"values": [
296+
{"string_value": "2"},
297+
{"string_value": "Bob"},
298+
{"string_value": "2023-10-02T00:00:00Z"},
299+
]
300+
},
301+
{
302+
"values": [
303+
{"string_value": "3"},
304+
{"string_value": "Charlie"},
305+
{"string_value": "2023-10-03T00:00:00Z"},
306+
]
307+
},
308+
{
309+
"values": [
310+
{"string_value": "4"},
311+
{"string_value": "David"},
312+
{"string_value": "2023-10-04T00:00:00Z"},
313+
]
314+
},
315+
{
316+
"values": [
317+
{"string_value": "5"},
318+
{"string_value": "Eve"},
319+
{"string_value": "2023-10-05T00:00:00Z"},
320+
]
321+
},
322+
{
323+
"values": [
324+
{"string_value": "6"},
325+
{"string_value": "Frank"},
326+
{"string_value": "2023-10-06T00:00:00Z"},
327+
]
328+
},
329+
{
330+
"values": [
331+
{"string_value": "7"},
332+
{"string_value": "Grace"},
333+
{"string_value": "2023-10-07T00:00:00Z"},
334+
]
335+
},
336+
{
337+
"values": [
338+
{"string_value": "8"},
339+
{"string_value": "Hank"},
340+
{"string_value": "2023-10-08T00:00:00Z"},
341+
]
342+
},
343+
{
344+
"values": [
345+
{"string_value": "9"},
346+
{"string_value": "Ivy"},
347+
{"string_value": "2023-10-09T00:00:00Z"},
348+
]
349+
},
350+
{
351+
"values": [
352+
{"string_value": "10"},
353+
{"string_value": "Jack"},
354+
{"string_value": "2023-10-10T00:00:00Z"},
355+
]
356+
},
285357
]
286358
},
287359
}
@@ -353,7 +425,7 @@ def markdown_to_dataframe(markdown_str: str) -> pd.DataFrame:
353425

354426
# Strip whitespace from column names and values
355427
df.columns = [col.strip() for col in df.columns]
356-
df = df.applymap(lambda x: x.strip() if isinstance(x, str) else x)
428+
df = df.map(lambda x: x.strip() if isinstance(x, str) else x)
357429

358430
# Drop the first column
359431
df = df.drop(columns=[df.columns[0]])
@@ -387,7 +459,11 @@ def test_parse_query_result_trims_large_data(max_tokens):
387459
"values": [
388460
{"string_value": str(i + 1)},
389461
{"string_value": random.choice(names)},
390-
{"string_value": (base_date + timedelta(days=random.randint(0, 365))).strftime("%Y-%m-%dT%H:%M:%SZ")},
462+
{
463+
"string_value": (
464+
base_date + timedelta(days=random.randint(0, 365))
465+
).strftime("%Y-%m-%dT%H:%M:%SZ")
466+
},
391467
]
392468
}
393469
for i in range(1000)
@@ -414,7 +490,8 @@ def test_parse_query_result_trims_large_data(max_tokens):
414490
"id": [int(row["values"][0]["string_value"]) for row in data_array],
415491
"name": [row["values"][1]["string_value"] for row in data_array],
416492
"created_at": [
417-
datetime.strptime(row["values"][2]["string_value"], "%Y-%m-%dT%H:%M:%SZ") for row in data_array
493+
datetime.strptime(row["values"][2]["string_value"], "%Y-%m-%dT%H:%M:%SZ")
494+
for row in data_array
418495
],
419496
}
420497
)

0 commit comments

Comments
 (0)