From 80bf6b3af73da7cc6f360963d669c7074bca8cef Mon Sep 17 00:00:00 2001
From: Cor Zuurmond <jczuurmond@protonmail.com>
Date: Tue, 19 Nov 2024 12:13:12 +0100
Subject: [PATCH] Limit integration test query

---
 tests/integration/test_core.py | 85 ++++++++++++++++++----------------
 1 file changed, 46 insertions(+), 39 deletions(-)

diff --git a/tests/integration/test_core.py b/tests/integration/test_core.py
index 4983a0d0..d9277263 100644
--- a/tests/integration/test_core.py
+++ b/tests/integration/test_core.py
@@ -17,49 +17,56 @@ def test_sql_execution_chunked(ws, disposition):
     assert total == 1999999000000
 
 
-def test_sql_execution(ws, env_or_skip):
-    results = []
+NYC_TAXI_LIMITED_TRIPS = """
+WITH zipcodes AS (
+  SELECT DISTINCT pickup_zip, dropoff_zip 
+  FROM samples.nyctaxi.trips 
+  WHERE pickup_zip = 10282 AND dropoff_zip <= 10005
+)
+
+SELECT 
+  trips.pickup_zip, 
+  trips.dropoff_zip, 
+  trips.tpep_pickup_datetime, 
+  trips.tpep_dropoff_datetime 
+FROM 
+    zipcodes 
+  JOIN 
+    samples.nyctaxi.trips AS trips 
+    ON zipcodes.pickup_zip = trips.pickup_zip AND zipcodes.dropoff_zip = trips.dropoff_zip
+ORDER BY trips.dropoff_zip, trips.tpep_pickup_datetime, trips.tpep_dropoff_datetime 
+"""
+
+
+def test_sql_execution(ws, env_or_skip) -> None:
+    results = set()
     see = StatementExecutionExt(ws, warehouse_id=env_or_skip("TEST_DEFAULT_WAREHOUSE_ID"))
-    for pickup_zip, dropoff_zip in see.fetch_all(
-        "SELECT pickup_zip, dropoff_zip FROM nyctaxi.trips LIMIT 10", catalog="samples"
-    ):
-        results.append((pickup_zip, dropoff_zip))
-    assert results == [
-        (10282, 10171),
-        (10110, 10110),
-        (10103, 10023),
-        (10022, 10017),
-        (10110, 10282),
-        (10009, 10065),
-        (10153, 10199),
-        (10112, 10069),
-        (10023, 10153),
-        (10012, 10003),
-    ]
-
-
-def test_sql_execution_partial(ws, env_or_skip):
-    results = []
+    for pickup_zip, dropoff_zip, *_ in see.fetch_all(NYC_TAXI_LIMITED_TRIPS, catalog="samples"):
+        results.add((pickup_zip, dropoff_zip))
+    assert results == {
+        (10282, 7114),
+        (10282, 10001),
+        (10282, 10002),
+        (10282, 10003),
+        (10282, 10005),
+    }
+
+
+def test_sql_execution_partial(ws, env_or_skip) -> None:
+    results = set()
     see = StatementExecutionExt(ws, warehouse_id=env_or_skip("TEST_DEFAULT_WAREHOUSE_ID"), catalog="samples")
-    for row in see("SELECT * FROM nyctaxi.trips LIMIT 10"):
-        pickup_time, dropoff_time = row[0], row[1]
-        pickup_zip = row.pickup_zip
-        dropoff_zip = row["dropoff_zip"]
+    for row in see(NYC_TAXI_LIMITED_TRIPS):
+        pickup_zip, dropoff_zip, pickup_time, dropoff_time = row[0], row[1], row[2], row[3]
         all_fields = row.asDict()
         logger.info(f"{pickup_zip}@{pickup_time} -> {dropoff_zip}@{dropoff_time}: {all_fields}")
-        results.append((pickup_zip, dropoff_zip))
-    assert results == [
-        (10282, 10171),
-        (10110, 10110),
-        (10103, 10023),
-        (10022, 10017),
-        (10110, 10282),
-        (10009, 10065),
-        (10153, 10199),
-        (10112, 10069),
-        (10023, 10153),
-        (10012, 10003),
-    ]
+        results.add((pickup_zip, dropoff_zip))
+    assert results == {
+        (10282, 7114),
+        (10282, 10001),
+        (10282, 10002),
+        (10282, 10003),
+        (10282, 10005),
+    }
 
 
 def test_fetch_one(ws):