From 899170679b44b8eebd5efe1af373eded98cbe9aa Mon Sep 17 00:00:00 2001
From: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com>
Date: Tue, 20 Jun 2023 16:21:49 -0700
Subject: [PATCH] Add command line arguments for testing queries (#1161)

* add filepath arg

* data_dir and queries_dir

---------

Co-authored-by: Ayush Dattagupta <ayushdg95@gmail.com>
---
 conftest.py                | 12 ++++++++++++
 tests/unit/test_queries.py | 28 ++++++++++++++++++----------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/conftest.py b/conftest.py
index d840bfd86..0042c9ca7 100644
--- a/conftest.py
+++ b/conftest.py
@@ -7,6 +7,8 @@
 def pytest_addoption(parser):
     parser.addoption("--rungpu", action="store_true", help="run tests meant for GPU")
     parser.addoption("--runqueries", action="store_true", help="run test queries")
+    parser.addoption("--data_dir", help="specify file path to the data")
+    parser.addoption("--queries_dir", help="specify file path to the queries")
 
 
 def pytest_runtest_setup(item):
@@ -21,3 +23,13 @@ def pytest_runtest_setup(item):
         dask.config.set({"dataframe.shuffle.algorithm": None})
     if "queries" in item.keywords and not item.config.getoption("--runqueries"):
         pytest.skip("need --runqueries option to run")
+
+
+@pytest.fixture(scope="session")
+def data_dir(request):
+    return request.config.getoption("--data_dir")
+
+
+@pytest.fixture(scope="session")
+def queries_dir(request):
+    return request.config.getoption("--queries_dir")
diff --git a/tests/unit/test_queries.py b/tests/unit/test_queries.py
index 42ef0c895..b32e9530f 100644
--- a/tests/unit/test_queries.py
+++ b/tests/unit/test_queries.py
@@ -51,15 +51,17 @@
 
 
 @pytest.fixture(scope="module")
-def c():
+def c(data_dir):
     # Lazy import, otherwise the pytest framework has problems
     from dask_sql.context import Context
 
     c = Context()
-    for table_name in os.listdir(f"{os.path.dirname(__file__)}/data/"):
+    if not data_dir:
+        data_dir = f"{os.path.dirname(__file__)}/data/"
+    for table_name in os.listdir(data_dir):
         c.create_table(
             table_name,
-            f"{os.path.dirname(__file__)}/data/{table_name}",
+            data_dir + "/" + table_name,
             format="parquet",
             gpu=False,
         )
@@ -68,17 +70,19 @@ def c():
 
 
 @pytest.fixture(scope="module")
-def gpu_c():
+def gpu_c(data_dir):
     pytest.importorskip("dask_cudf")
 
     # Lazy import, otherwise the pytest framework has problems
     from dask_sql.context import Context
 
     c = Context()
-    for table_name in os.listdir(f"{os.path.dirname(__file__)}/data/"):
+    if not data_dir:
+        data_dir = f"{os.path.dirname(__file__)}/data/"
+    for table_name in os.listdir(data_dir):
         c.create_table(
             table_name,
-            f"{os.path.dirname(__file__)}/data/{table_name}",
+            data_dir + "/" + table_name,
             format="parquet",
             gpu=True,
         )
@@ -88,8 +92,10 @@ def gpu_c():
 
 @pytest.mark.queries
 @pytest.mark.parametrize("query", QUERIES)
-def test_query(c, client, query):
-    with open(f"{os.path.dirname(__file__)}/queries/{query}") as f:
+def test_query(c, client, query, queries_dir):
+    if not queries_dir:
+        queries_dir = f"{os.path.dirname(__file__)}/queries/"
+    with open(queries_dir + "/" + query) as f:
         sql = f.read()
 
     res = c.sql(sql)
@@ -99,8 +105,10 @@ def test_query(c, client, query):
 @pytest.mark.gpu
 @pytest.mark.queries
 @pytest.mark.parametrize("query", QUERIES)
-def test_gpu_query(gpu_c, gpu_client, query):
-    with open(f"{os.path.dirname(__file__)}/queries/{query}") as f:
+def test_gpu_query(gpu_c, gpu_client, query, queries_dir):
+    if not queries_dir:
+        queries_dir = f"{os.path.dirname(__file__)}/queries/"
+    with open(queries_dir + "/" + query) as f:
         sql = f.read()
 
     res = gpu_c.sql(sql)