add connection unit tests

sfu-db · Apr 19, 2024 · 9bc3495 · 9bc3495
1 parent de64529
commit 9bc3495
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 21 deletions.
diff --git a/connectorx-python/connectorx/tests/test_bigquery.py b/connectorx-python/connectorx/tests/test_bigquery.py
@@ -4,7 +4,7 @@
 import pytest
 from pandas.testing import assert_frame_equal
 
-from .. import read_sql
+from .. import read_sql, Connection
 
 
 @pytest.fixture(scope="module")  # type: ignore
@@ -121,9 +121,7 @@ def test_bigquery_some_empty_partition(bigquery_url: str) -> None:
         index=range(1),
         data={
             "test_int": pd.Series([1], dtype="Int64"),
-            "test_string": pd.Series(
-                ["str1"], dtype="object"
-            ),
+            "test_string": pd.Series(["str1"], dtype="object"),
             "test_float": pd.Series([1.10], dtype="float64"),
             "test_bool": pd.Series([True], dtype="boolean"),
         },
@@ -137,10 +135,7 @@ def test_bigquery_some_empty_partition(bigquery_url: str) -> None:
 )
 def test_bigquery_join(bigquery_url: str) -> None:
     query = "SELECT T.test_int, T.test_string, S.test_str FROM `dataprep-bigquery.dataprep.test_table` T INNER JOIN `dataprep-bigquery.dataprep.test_types` S ON T.test_int = S.test_int"
-    df = read_sql(
-        bigquery_url,
-        query
-    )
+    df = read_sql(bigquery_url, query)
     df = df.sort_values("test_int").reset_index(drop=True)
     expected = pd.DataFrame(
         index=range(2),
@@ -151,14 +146,14 @@ def test_bigquery_join(bigquery_url: str) -> None:
                     "str1",
                     "str2",
                 ],
-                dtype="object"
+                dtype="object",
             ),
             "test_str": pd.Series(
                 [
                     "😁😂😜",
                     "こんにちはЗдра́в",
                 ],
-                dtype="object"
+                dtype="object",
             ),
         },
     )
@@ -188,22 +183,21 @@ def test_bigquery_join_with_partition(bigquery_url: str) -> None:
                     "str1",
                     "str2",
                 ],
-                dtype="object"
+                dtype="object",
             ),
             "test_str": pd.Series(
                 [
                     "😁😂😜",
                     "こんにちはЗдра́в",
                 ],
-                dtype="object"
+                dtype="object",
             ),
         },
     )
     df.sort_values(by="test_int", inplace=True, ignore_index=True)
     assert_frame_equal(df, expected, check_names=True)
 
 
-
 @pytest.mark.skipif(
     not os.environ.get("BIGQUERY_URL"),
     reason="Test bigquery only when `BIGQUERY_URL` is set",
@@ -310,3 +304,11 @@ def test_bigquery_types(bigquery_url: str) -> None:
         },
     )
     assert_frame_equal(df, expected, check_names=True)
+
+
+@pytest.mark.skipif(
+    not os.environ.get("BIGQUERY_URL"),
+    reason="Test bigquery only when `BIGQUERY_URL` is set",
+)
+def test_connection(bigquery_url: str) -> None:
+    test_bigquery_types(Connection(bigquery_url))
diff --git a/connectorx-python/connectorx/tests/test_clickhouse.py b/connectorx-python/connectorx/tests/test_clickhouse.py
@@ -4,7 +4,7 @@
 import pytest
 from pandas.testing import assert_frame_equal
 
-from .. import read_sql
+from .. import read_sql, Connection
 
 
 @pytest.fixture(scope="module")  # type: ignore
@@ -81,3 +81,11 @@ def test_clickhouse_types(clickhouse_url: str) -> None:
         },
     )
     assert_frame_equal(df, expected, check_names=True)
+
+
+@pytest.mark.skipif(
+    not os.environ.get("CLICKHOUSE_URL"),
+    reason="Do not test Clickhouse unless `CLICKHOUSE_URL` is set",
+)
+def test_connection(clickhouse_url: str) -> None:
+    test_clickhouse_types(Connection(clickhouse_url))
diff --git a/connectorx-python/connectorx/tests/test_mssql.py b/connectorx-python/connectorx/tests/test_mssql.py
@@ -3,6 +3,7 @@
 import pandas as pd
 import pytest
 from pandas.testing import assert_frame_equal
+from connectorx import Connection
 
 from .. import read_sql
 
@@ -92,7 +93,6 @@ def test_mssql_udf(mssql_url: str) -> None:
 
 
 def test_manual_partition(mssql_url: str) -> None:
-
     queries = [
         "SELECT * FROM test_table WHERE test_int < 2",
         "SELECT * FROM test_table WHERE test_int >= 2",
@@ -496,3 +496,7 @@ def test_mssql_offset(mssql_url: str) -> None:
         }
     )
     assert_frame_equal(df, expected, check_names=True)
+
+
+def test_connection(mssql_url: str) -> None:
+    test_mssql_offset(Connection(mssql_url))
diff --git a/connectorx-python/connectorx/tests/test_mysql.py b/connectorx-python/connectorx/tests/test_mysql.py
@@ -4,7 +4,7 @@
 import pytest
 from pandas.testing import assert_frame_equal
 
-from .. import read_sql
+from .. import read_sql, Connection
 
 
 @pytest.fixture(scope="module")  # type: ignore
@@ -468,3 +468,7 @@ def test_mysql_cte(mysql_url: str) -> None:
         },
     )
     assert_frame_equal(df, expected, check_names=True)
+
+
+def test_connection(mysql_url: str) -> None:
+    test_mysql_cte(Connection(mysql_url))
diff --git a/connectorx-python/connectorx/tests/test_oracle.py b/connectorx-python/connectorx/tests/test_oracle.py
@@ -4,7 +4,7 @@
 import pytest
 from pandas.testing import assert_frame_equal
 
-from .. import read_sql
+from .. import read_sql, Connection
 
 
 @pytest.fixture(scope="module")  # type: ignore
@@ -440,4 +440,11 @@ def test_oracle_round_function(oracle_url: str) -> None:
             "TEST_ROUND": pd.Series([1.11, 2.22, 3.33, None], dtype="float64"),
         }
     )
-    assert_frame_equal(df, expected, check_names=True)
+    assert_frame_equal(df, expected, check_names=True)
+
+
+@pytest.mark.skipif(
+    not os.environ.get("ORACLE_URL"), reason="Test oracle only when `ORACLE_URL` is set"
+)
+def test_connection(oracle_url: str) -> None:
+    test_oracle_round_function(Connection(oracle_url))
diff --git a/connectorx-python/connectorx/tests/test_redshift.py b/connectorx-python/connectorx/tests/test_redshift.py
@@ -5,7 +5,7 @@
 import pytest
 from pandas.testing import assert_frame_equal
 
-from .. import read_sql
+from .. import read_sql, Connection
 
 
 @pytest.fixture(scope="module")  # type: ignore
@@ -134,3 +134,11 @@ def test_read_sql_on_utf8(redshift_url: str) -> None:
         },
     )
     assert_frame_equal(df, expected, check_names=True)
+
+
+@pytest.mark.skipif(
+    not os.environ.get("REDSHIFT_URL"),
+    reason="Do not test Redshift unless `REDSHIFT_URL` is set",
+)
+def test_connection(redshift_url: str) -> None:
+    test_read_sql_on_utf8(Connection(redshift_url))
diff --git a/connectorx-python/connectorx/tests/test_sqlite.py b/connectorx-python/connectorx/tests/test_sqlite.py
@@ -5,7 +5,7 @@
 import pytest
 from pandas.testing import assert_frame_equal
 
-from .. import read_sql
+from .. import read_sql, Connection
 
 
 @pytest.fixture(scope="module")  # type: ignore
@@ -215,7 +215,6 @@ def test_sqlite_with_partition(sqlite_db: str) -> None:
 
 
 def test_manual_partition(sqlite_db: str) -> None:
-
     queries = [
         "SELECT test_int, test_nullint, test_str, test_float, test_bool, test_date, test_time, test_datetime FROM test_table WHERE test_int < 2",
         "SELECT test_int, test_nullint, test_str, test_float, test_bool, test_date, test_time, test_datetime FROM test_table WHERE test_int >= 2",
@@ -391,3 +390,7 @@ def test_sqlite_cte(sqlite_db: str) -> None:
         },
     )
     assert_frame_equal(df, expected, check_names=True)
+
+
+def test_connection(sqlite_db: str) -> None:
+    test_sqlite_cte(Connection(sqlite_db))