chore: cleanup in python bindings tests

bindings/python/tests/test_connection.py

@@ -3,25 +3,28 @@
 import pandas as pd
 
 
+def test_default_conn_uses_same_db():
+    # Create table
+    glaredb.execute("create table other (a int)")
+    # Try to query it. This would error if we weren't using the same db.
+    glaredb.execute("select * from other")
+    glaredb.execute("drop table other")
+
+
 def test_with_context():
     with glaredb.connect() as con:
         con.execute("select 1")
 
 
 def test_default_connection_sql():
-    out = glaredb.sql("select 1 as a").to_pandas()
-    expected = pd.DataFrame({"a": [1]})
-    assert out.equals(expected)
+    with glaredb.connect() as db:
+        out = db.sql("select 1 as a").to_pandas()
+        expected = pd.DataFrame({"a": [1]})
+        assert out.equals(expected)
 
 
 def test_default_connection_execute():
-    out = glaredb.execute("select 1 as a").to_pandas()
-    expected = pd.DataFrame({"a": [1]})
-    assert out.equals(expected)
-
-
-# def test_default_conn_uses_same_db():
-#     # Create table
-#     glaredb.execute("create table hello (a int)")
-#     # Try to query it. This would error if we weren't using the same db.
-#     glaredb.execute("select * from hello")
+    with glaredb.connect() as db:
+        out = db.execute("select 1 as a").to_pandas()
+        expected = pd.DataFrame({"a": [1]})
+        assert out.equals(expected)

bindings/python/tests/test_logical_plan.py

@@ -3,16 +3,15 @@
 
 
 def test_compose():
-    con = glaredb.connect()
-    con.execute("create table tbl (a int, b int);")
-    con.execute("insert into tbl values (1, 2);")
-    con.execute("insert into tbl values (3, 4);")
-    con.execute("insert into tbl values (5, 6);")
-    intermediate = con.sql("select * from tbl where a > 2 order by a;")
-    out_1 = intermediate.to_arrow().to_pydict()
-    expected = {"a": [3, 5], "b": [4, 6]}
-    assert out_1 == expected
-    out_2 = con.sql("select * from intermediate where b > 4;").to_arrow().to_pydict()
-    expected = {"a": [5], "b": [6]}
-    assert out_2 == expected
-    con.close()
+    with glaredb.connect() as con:
+        con.execute("create table tbl (a int, b int);")
+        con.execute("insert into tbl (a, b) values (1, 2);")
+        con.execute("insert into tbl values (3, 4);")
+        con.execute("insert into tbl values (5, 6);")
+        intermediate = con.sql("select * from tbl where a > 2 order by a;")
+        out_1 = intermediate.to_arrow().to_pydict()
+        expected = {"a": [3, 5], "b": [4, 6]}
+        assert out_1 == expected
+        out_2 = con.sql("select * from intermediate where b > 4;").to_arrow().to_pydict()
+        expected = {"a": [5], "b": [6]}
+        assert out_2 == expected