orca-sdk 0.0.96__py3-none-any.whl → 0.0.98__py3-none-any.whl

orca_sdk/datasource_test.py

@@ -2,6 +2,8 @@ import os
 import tempfile
 from uuid import uuid4
 
+import pandas as pd
+import pyarrow as pa
 import pytest
 
 from .datasource import Datasource
@@ -102,3 +104,195 @@ def test_download_datasource(datasource):
         output_path = os.path.join(temp_dir, "datasource.zip")
         datasource.download(output_path)
         assert os.path.exists(output_path)
+
+
+def test_from_list():
+    # Test creating datasource from list of dictionaries
+    data = [
+        {"column1": 1, "column2": "a"},
+        {"column1": 2, "column2": "b"},
+        {"column1": 3, "column2": "c"},
+    ]
+    datasource = Datasource.from_list(f"test_list_{uuid4()}", data)
+    assert datasource.name.startswith("test_list_")
+    assert datasource.length == 3
+    assert "column1" in datasource.columns
+    assert "column2" in datasource.columns
+
+
+def test_from_dict():
+    # Test creating datasource from dictionary of columns
+    data = {
+        "column1": [1, 2, 3],
+        "column2": ["a", "b", "c"],
+    }
+    datasource = Datasource.from_dict(f"test_dict_{uuid4()}", data)
+    assert datasource.name.startswith("test_dict_")
+    assert datasource.length == 3
+    assert "column1" in datasource.columns
+    assert "column2" in datasource.columns
+
+
+def test_from_pandas():
+    # Test creating datasource from pandas DataFrame
+    df = pd.DataFrame(
+        {
+            "column1": [1, 2, 3],
+            "column2": ["a", "b", "c"],
+        }
+    )
+    datasource = Datasource.from_pandas(f"test_pandas_{uuid4()}", df)
+    assert datasource.name.startswith("test_pandas_")
+    assert datasource.length == 3
+    assert "column1" in datasource.columns
+    assert "column2" in datasource.columns
+
+
+def test_from_arrow():
+    # Test creating datasource from pyarrow Table
+    table = pa.table(
+        {
+            "column1": [1, 2, 3],
+            "column2": ["a", "b", "c"],
+        }
+    )
+    datasource = Datasource.from_arrow(f"test_arrow_{uuid4()}", table)
+    assert datasource.name.startswith("test_arrow_")
+    assert datasource.length == 3
+    assert "column1" in datasource.columns
+    assert "column2" in datasource.columns
+
+
+def test_from_list_already_exists():
+    # Test the if_exists parameter with from_list
+    data = [{"column1": 1, "column2": "a"}]
+    name = f"test_list_exists_{uuid4()}"
+
+    # Create the first datasource
+    datasource1 = Datasource.from_list(name, data)
+    assert datasource1.length == 1
+
+    # Try to create again with if_exists="error" (should raise)
+    with pytest.raises(ValueError):
+        Datasource.from_list(name, data, if_exists="error")
+
+    # Try to create again with if_exists="open" (should return existing)
+    datasource2 = Datasource.from_list(name, data, if_exists="open")
+    assert datasource2.id == datasource1.id
+    assert datasource2.name == datasource1.name
+
+
+def test_from_dict_already_exists():
+    # Test the if_exists parameter with from_dict
+    data = {"column1": [1], "column2": ["a"]}
+    name = f"test_dict_exists_{uuid4()}"
+
+    # Create the first datasource
+    datasource1 = Datasource.from_dict(name, data)
+    assert datasource1.length == 1
+
+    # Try to create again with if_exists="error" (should raise)
+    with pytest.raises(ValueError):
+        Datasource.from_dict(name, data, if_exists="error")
+
+    # Try to create again with if_exists="open" (should return existing)
+    datasource2 = Datasource.from_dict(name, data, if_exists="open")
+    assert datasource2.id == datasource1.id
+    assert datasource2.name == datasource1.name
+
+
+def test_from_pandas_already_exists():
+    # Test the if_exists parameter with from_pandas
+    df = pd.DataFrame({"column1": [1], "column2": ["a"]})
+    name = f"test_pandas_exists_{uuid4()}"
+
+    # Create the first datasource
+    datasource1 = Datasource.from_pandas(name, df)
+    assert datasource1.length == 1
+
+    # Try to create again with if_exists="error" (should raise)
+    with pytest.raises(ValueError):
+        Datasource.from_pandas(name, df, if_exists="error")
+
+    # Try to create again with if_exists="open" (should return existing)
+    datasource2 = Datasource.from_pandas(name, df, if_exists="open")
+    assert datasource2.id == datasource1.id
+    assert datasource2.name == datasource1.name
+
+
+def test_from_arrow_already_exists():
+    # Test the if_exists parameter with from_arrow
+    table = pa.table({"column1": [1], "column2": ["a"]})
+    name = f"test_arrow_exists_{uuid4()}"
+
+    # Create the first datasource
+    datasource1 = Datasource.from_arrow(name, table)
+    assert datasource1.length == 1
+
+    # Try to create again with if_exists="error" (should raise)
+    with pytest.raises(ValueError):
+        Datasource.from_arrow(name, table, if_exists="error")
+
+    # Try to create again with if_exists="open" (should return existing)
+    datasource2 = Datasource.from_arrow(name, table, if_exists="open")
+    assert datasource2.id == datasource1.id
+    assert datasource2.name == datasource1.name
+
+
+def test_from_disk_csv():
+    # Test creating datasource from CSV file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+        f.write("column1,column2\n1,a\n2,b\n3,c")
+        f.flush()
+
+        try:
+            datasource = Datasource.from_disk(f"test_csv_{uuid4()}", f.name)
+            assert datasource.length == 3
+            assert "column1" in datasource.columns
+            assert "column2" in datasource.columns
+        finally:
+            os.unlink(f.name)
+
+
+def test_from_disk_json():
+    # Test creating datasource from JSON file
+    import json
+
+    data = [{"column1": 1, "column2": "a"}, {"column1": 2, "column2": "b"}]
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(data, f)
+        f.flush()
+
+        try:
+            datasource = Datasource.from_disk(f"test_json_{uuid4()}", f.name)
+            assert datasource.length == 2
+            assert "column1" in datasource.columns
+            assert "column2" in datasource.columns
+        finally:
+            os.unlink(f.name)
+
+
+def test_from_disk_already_exists():
+    # Test the if_exists parameter with from_disk
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+        f.write("column1,column2\n1,a")
+        f.flush()
+
+        try:
+            name = f"test_disk_exists_{uuid4()}"
+
+            # Create the first datasource
+            datasource1 = Datasource.from_disk(name, f.name)
+            assert datasource1.length == 1
+
+            # Try to create again with if_exists="error" (should raise)
+            with pytest.raises(ValueError):
+                Datasource.from_disk(name, f.name, if_exists="error")
+
+            # Try to create again with if_exists="open" (should return existing)
+            datasource2 = Datasource.from_disk(name, f.name, if_exists="open")
+            assert datasource2.id == datasource1.id
+            assert datasource2.name == datasource1.name
+        finally:
+            os.unlink(f.name)