Skip to content

Commit 51cef2b

Browse files
validate row count for sync query tests as well
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent f9f9f31 commit 51cef2b

File tree

1 file changed

+21
-2
lines changed

1 file changed

+21
-2
lines changed

examples/experimental/tests/test_sea_sync_query.py

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,13 @@ def test_sea_sync_query_with_cloud_fetch():
7272
f"{actual_row_count} rows retrieved against {requested_row_count} requested"
7373
)
7474

75+
# Verify total row count
76+
if actual_row_count != requested_row_count:
77+
logger.error(
78+
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
79+
)
80+
return False
81+
7582
# Close resources
7683
cursor.close()
7784
connection.close()
@@ -132,15 +139,27 @@ def test_sea_sync_query_without_cloud_fetch():
132139
# For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
133140
requested_row_count = 100
134141
cursor = connection.cursor()
135-
logger.info("Executing synchronous query without cloud fetch: SELECT 100 rows")
142+
logger.info(
143+
f"Executing synchronous query without cloud fetch: SELECT {requested_row_count} rows"
144+
)
136145
cursor.execute(
137146
"SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
138147
)
139148

140149
results = [cursor.fetchone()]
141150
results.extend(cursor.fetchmany(10))
142151
results.extend(cursor.fetchall())
143-
logger.info(f"{len(results)} rows retrieved against 100 requested")
152+
actual_row_count = len(results)
153+
logger.info(
154+
f"{actual_row_count} rows retrieved against {requested_row_count} requested"
155+
)
156+
157+
# Verify total row count
158+
if actual_row_count != requested_row_count:
159+
logger.error(
160+
f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
161+
)
162+
return False
144163

145164
# Close resources
146165
cursor.close()

0 commit comments

Comments
 (0)