chore: pytest hacking (#2606)

justfile

@@ -8,9 +8,8 @@ alias slt := sql-logic-tests
 
 os_arch := os() + '-' + arch()
 
-
+VENV := ".venv"
 VENV_BIN := VENV / "bin"
-VENV := env_var_or_default("VENV", ".venv")
 
 # Run benchmarks subcommands. see `benchmarks/justfile` for more details.
 bench cmd *args:
@@ -122,11 +121,10 @@ protoc:
 # Installs python dependencies for testing
 venv:
   python3 -c "import virtualenv" || python3 -m pip --quiet install virtualenv
-  python3 -m virtualenv {{VENV}} --quiet
+  python3 -m virtualenv .venv --quiet
   {{VENV_BIN}}/python -m pip install poetry
   {{VENV_BIN}}/poetry -C tests install
 
-
 # Runs pytest in the tests directory.
 pytest *args:
   {{VENV_BIN}}/poetry -C tests lock --no-update

tests/tests/test_bson.py → tests/test_bson.py

@@ -9,8 +9,8 @@
 import psycopg2.extras
 import pytest
 
-from fixtures.glaredb import glaredb_connection, debug_path
-import tools
+from tests.fixtures.glaredb import glaredb_connection, debug_path
+import tests.tools
 
 
 def test_bson_copy_to(
@@ -45,7 +45,7 @@ def test_bson_copy_to(
             assert "amount" in doc
             assert doc["amount"] == idx
 
-    with tools.cd(output_dir):
+    with tests.tools.cd(output_dir):
         out = subprocess.run(
             [
                 f"{debug_path}",

tests/tests/test_dbt.py → tests/test_dbt.py

@@ -1,36 +1,34 @@
 import pathlib
 import tests
+import os
 
 import psycopg2.extensions
+from dbt.cli.main import dbtRunner, dbtRunnerResult
 import pytest
-import os
 
+import tests.tools
 from tests.fixtures.glaredb import glaredb_connection, debug_path
 
-from dbt.cli.main import dbtRunner, dbtRunnerResult
 
 @pytest.fixture
 def dbt_project_path() -> pathlib.Path:
     return tests.PKG_DIRECTORY.joinpath("tests", "fixtures", "dbt_project")
 
-@pytest.mark.parametrize("model_name,run_success,query_result",
-                         [
-                             ("table_materialization", True, 10),
-                             pytest.param("view_materialization", True, 10, marks=pytest.mark.xfail),
-                         ]
-                         )
+
+@pytest.mark.parametrize(
+    "model_name,run_success,query_result",
+    [
+        ("table_materialization", True, 10),
+        pytest.param("view_materialization", True, 10, marks=pytest.mark.xfail),
+    ],
+)
 def test_dbt_glaredb(
     glaredb_connection: psycopg2.extensions.connection,
     dbt_project_path,
     model_name,
     run_success,
-    query_result
+    query_result,
 ):
-    print("LISTDIR", os.listdir('.'))
-    dbt: dbtRunner = dbtRunner()
-
-    os.environ["DBT_USER"] = glaredb_connection.info.user
-
     dbt_project_directory: pathlib.Path = dbt_project_path
     dbt_profiles_directory: pathlib.Path = dbt_project_path
 
@@ -40,22 +38,23 @@ def test_dbt_glaredb(
             "INSERT INTO dbt_test (amount) VALUES (0), (1), (2), (3), (4), (5), (6), (7), (8), (9)"
         )
 
-    cli_args: list = [
-        "run",
-        "--project-dir",
-        dbt_project_directory,
-        "--profiles-dir",
-        dbt_profiles_directory,
-        "-m",
-        model_name
-    ]
-    #
-    res: dbtRunnerResult = dbt.invoke(cli_args)
+    with tests.tools.env("DBT_USER", glaredb_connection.info.user):
+        res: dbtRunnerResult = dbtRunner().invoke(
+            [
+                "run",
+                "--project-dir",
+                dbt_project_directory,
+                "--profiles-dir",
+                dbt_profiles_directory,
+                "-m",
+                model_name,
+            ]
+        )
 
-    assert res.success is run_success
+        assert res.success is run_success
 
     with glaredb_connection.cursor() as curr:
         curr.execute(f"select count(*) from {model_name}")
         result: list = curr.fetchone()[0]
 
-    assert result == query_result
+        assert result == query_result

tests/tests/test_functions.py → tests/test_functions.py

@@ -1,6 +1,6 @@
 import psycopg2
 
-from fixtures.glaredb import glaredb_connection, debug_path
+from tests.fixtures.glaredb import glaredb_connection, debug_path
 
 
 def test_scalar_parsing(

tests/tests/test_json.py → tests/test_json.py

@@ -6,7 +6,8 @@
 import psycopg2.extras
 import pytest
 
-from fixtures.glaredb import glaredb_connection, debug_path
+from tests.fixtures.glaredb import glaredb_connection, debug_path
+
 
 @pytest.fixture
 def beatle_mock_data():
@@ -21,7 +22,7 @@ def beatle_mock_data():
             "case": i + 1,
             "rand": random.random(),
         }
-        
+
         if beatle_id == 0:
             obj["house"] = "the dakota"
         if beatle_id == 1:
@@ -39,27 +40,25 @@ def beatle_mock_data():
 def test_read_json_array(
     glaredb_connection: psycopg2.extensions.connection,
     tmp_path_factory: pytest.TempPathFactory,
-    beatle_mock_data: list[dict], 
+    beatle_mock_data: list[dict],
 ):
     tmp_dir = tmp_path_factory.mktemp(basename="read-json-array-", numbered=True)
     data_path = tmp_dir.joinpath("beatles.100.json")
 
     with open(data_path, "w") as f:
         json.dump(beatle_mock_data, f, indent="  ")
-            
+
     with glaredb_connection.cursor() as curr:
         curr.execute(f"select count(*) from read_json('{data_path}');")
         r = curr.fetchone()
         assert r[0] == 100
 
-    with glaredb_connection.cursor(
-            cursor_factory=psycopg2.extras.RealDictCursor
-    ) as curr:
+    with glaredb_connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curr:
         curr.execute(f"select * from read_json('{data_path}');")
         rows = curr.fetchall()
         assert len(rows) == 100
         for row in rows:
-            assert len(row) == 8 # superset schema
+            assert len(row) == 8  # superset schema
             assert "house" in row
             assert "beatle_name" in row
             if row["beatle_name"] == "john":
@@ -72,7 +71,7 @@ def test_read_json_array(
 def test_read_json_glob(
     glaredb_connection: psycopg2.extensions.connection,
     tmp_path_factory: pytest.TempPathFactory,
-    beatle_mock_data: list[dict], 
+    beatle_mock_data: list[dict],
 ):
     beatles = ["john", "paul", "george", "ringo"]
 
@@ -81,20 +80,18 @@ def test_read_json_glob(
     for idx, doc in enumerate(beatle_mock_data):
         with open(tmp_dir.joinpath(f"file-{idx}.json"), "w") as f:
             json.dump(doc, f, indent="  ")
-            
+
     with glaredb_connection.cursor() as curr:
         curr.execute(f"select count(*) from read_json('{tmp_dir}/*.json')")
         r = curr.fetchone()
         assert r[0] == 100
 
-    with glaredb_connection.cursor(
-            cursor_factory=psycopg2.extras.RealDictCursor
-    ) as curr:
+    with glaredb_connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as curr:
         curr.execute(f"select * from read_json('{data_path}');")
         rows = curr.fetchall()
         assert len(rows) == 100
         for row in rows:
-            assert len(row) == 8 # superset schema
+            assert len(row) == 8  # superset schema
             assert "house" in row
             assert "beatle_name" in row
             if row["beatle_name"] == "john":

tests/tests/test_lance.py → tests/test_lance.py

@@ -8,8 +8,8 @@
 import pytest
 
 
-import tools
-from fixtures.glaredb import glaredb_connection, debug_path
+import tests.tools
+from tests.fixtures.glaredb import glaredb_connection, debug_path
 
 
 def test_sanity_check(
@@ -64,7 +64,7 @@ def test_copy_to_round_trip(
 
     output_path_rel = tmp_path_factory.mktemp("lance-rel")
 
-    with tools.cd(output_path_rel):
+    with tests.tools.cd(output_path_rel):
         with glaredb_connection.cursor() as curr:
             curr.execute("COPY lance_test TO './' FORMAT lance")
 
@@ -91,10 +91,10 @@ def test_copy_to_round_trip(
         curr.execute(f"COPY lance_test TO '{output_path}' FORMAT lance")
         curr.execute(f"create external table lance_import from lance options (location '{output_path}')")
         curr.execute("alter table lance_import set access_mode to read_write")
-        
+
         for i in range(10):
             curr.execute("insert into lance_import values (%s)", str(i))
-        
+
     with glaredb_connection.cursor() as curr:
         curr.execute("select count(*) from lance_import;")
         assert curr.fetchone()[0] == 20

tests/tests/test_smoke.py → tests/test_smoke.py

@@ -5,7 +5,7 @@
 import pytest
 import psycopg2
 
-from fixtures.glaredb import glaredb_connection, release_path, debug_path
+from tests.fixtures.glaredb import glaredb_connection, release_path, debug_path
 
 
 def test_release_exists(release_path: pathlib.Path):

tests/test_tools.py

@@ -0,0 +1,23 @@
+import os
+import pytest
+
+import tests.tools
+
+
+def test_cd():
+    cwd = os.getcwd()
+
+    with tests.tools.cd("/tmp"):
+        assert not cwd == os.getcwd()
+
+    assert cwd == os.getcwd()
+
+
+def test_env():
+    assert not "merlin" in os.environ
+
+    with tests.tools.env("merlin", "cat"):
+        assert "merlin" in os.environ
+        assert os.environ["merlin"] == "cat"
+
+    assert not "merlin" in os.environ

tests/tools.py

@@ -1,7 +1,9 @@
+import collections.abc
 import contextlib
 import os
 import pathlib
 
+
 @contextlib.contextmanager
 def cd(path: pathlib.Path):
     cur = os.getcwd()
@@ -12,3 +14,18 @@ def cd(path: pathlib.Path):
         yield
     finally:
         os.chdir(cur)
+
+
+@contextlib.contextmanager
+def env(key: str, val: str):
+    prev = os.getenv(key)
+
+    os.environ[key] = val
+
+    try:
+        yield
+    finally:
+        if prev is None:
+            del os.environ[key]
+        else:
+            os.environ[key] = prev