Test

huggingface · Oct 7, 2022 · 353bec3 · 353bec3 · github-actions · Oct 7, 2022
1 parent c7a1961
commit 353bec3
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
@@ -3356,6 +3356,25 @@ def _check_sql_dataset(dataset, expected_features):
         assert dataset.features[feature].dtype == expected_dtype
 
 
+@require_sqlalchemy
+@pytest.mark.parametrize("con_type", ["string", "engine"])
+def test_dataset_from_sql_con_type(con_type, sqlite_path, tmp_path):
+    cache_dir = tmp_path / "cache"
+    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
+    if con_type == "string":
+        con = "sqlite:///" + sqlite_path
+    elif con_type == "engine":
+        import sqlalchemy
+
+        con = sqlalchemy.create_engine("sqlite:///" + sqlite_path)
+    dataset = Dataset.from_sql(
+        "dataset",
+        con,
+        cache_dir=cache_dir,
+    )
+    _check_sql_dataset(dataset, expected_features)
+
+
 @require_sqlalchemy
 @pytest.mark.parametrize(
     "features",