orca-sdk 0.1.9__py3-none-any.whl

orca_sdk/datasource_test.py

@@ -0,0 +1,463 @@
+import json
+import os
+import tempfile
+from typing import cast
+from uuid import uuid4
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+from datasets import Dataset
+
+from .datasource import Datasource
+
+
+def test_create_datasource(datasource, hf_dataset):
+    assert datasource is not None
+    assert datasource.name == "test_datasource"
+    assert datasource.length == len(hf_dataset)
+
+
+def test_create_datasource_unauthenticated(unauthenticated_client, hf_dataset):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            Datasource.from_hf_dataset("test_datasource", hf_dataset)
+
+
+def test_create_datasource_already_exists_error(hf_dataset, datasource):
+    with pytest.raises(ValueError):
+        Datasource.from_hf_dataset("test_datasource", hf_dataset, if_exists="error")
+
+
+def test_create_datasource_already_exists_return(hf_dataset, datasource):
+    returned_dataset = Datasource.from_hf_dataset("test_datasource", hf_dataset, if_exists="open")
+    assert returned_dataset is not None
+    assert returned_dataset.name == "test_datasource"
+    assert returned_dataset.length == len(hf_dataset)
+
+
+def test_open_datasource(datasource):
+    fetched_datasource = Datasource.open(datasource.name)
+    assert fetched_datasource is not None
+    assert fetched_datasource.name == datasource.name
+    assert fetched_datasource.length == len(datasource)
+
+
+def test_open_datasource_unauthenticated(unauthenticated_client, datasource):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            Datasource.open("test_datasource")
+
+
+def test_open_datasource_invalid_input():
+    with pytest.raises(ValueError, match=r"Invalid input:.*"):
+        Datasource.open("not valid id")
+
+
+def test_open_datasource_not_found():
+    with pytest.raises(LookupError):
+        Datasource.open(str(uuid4()))
+
+
+def test_open_datasource_unauthorized(unauthorized_client, datasource):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            Datasource.open(datasource.id)
+
+
+def test_all_datasources(datasource):
+    datasources = Datasource.all()
+    assert len(datasources) > 0
+    assert any(datasource.name == datasource.name for datasource in datasources)
+
+
+def test_all_datasources_unauthenticated(unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            Datasource.all()
+
+
+def test_drop_datasource(hf_dataset):
+    Datasource.from_hf_dataset("datasource_to_delete", hf_dataset)
+    assert Datasource.exists("datasource_to_delete")
+    Datasource.drop("datasource_to_delete")
+    assert not Datasource.exists("datasource_to_delete")
+
+
+def test_drop_datasource_unauthenticated(datasource, unauthenticated_client):
+    with unauthenticated_client.use():
+        with pytest.raises(ValueError, match="Invalid API key"):
+            Datasource.drop(datasource.id)
+
+
+def test_drop_datasource_not_found():
+    with pytest.raises(LookupError):
+        Datasource.drop(str(uuid4()))
+    # ignores error if specified
+    Datasource.drop(str(uuid4()), if_not_exists="ignore")
+
+
+def test_drop_datasource_unauthorized(datasource, unauthorized_client):
+    with unauthorized_client.use():
+        with pytest.raises(LookupError):
+            Datasource.drop(datasource.id)
+
+
+def test_drop_datasource_invalid_input():
+    with pytest.raises(ValueError, match=r"Invalid input:.*"):
+        Datasource.drop("not valid id")
+
+
+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)
+
+
+def test_query_datasource_rows():
+    """Test querying rows from a datasource with pagination and shuffle."""
+    # Create a new dataset with 5 entries for testing
+    test_data = [{"id": i, "name": f"item_{i}"} for i in range(5)]
+    datasource = Datasource.from_list(name="test_query_datasource", data=test_data)
+
+    # Test basic query
+    rows = datasource.query(limit=3)
+    assert len(rows) == 3
+    assert all(isinstance(row, dict) for row in rows)
+
+    # Test offset
+    offset_rows = datasource.query(offset=2, limit=2)
+    assert len(offset_rows) == 2
+    assert offset_rows[0]["id"] == 2
+
+    # Test shuffle
+    shuffled_rows = datasource.query(limit=5, shuffle=True)
+    assert len(shuffled_rows) == 5
+    assert not all(row["id"] == i for i, row in enumerate(shuffled_rows))
+
+    # Test shuffle with seed
+    assert datasource.query(limit=5, shuffle=True, shuffle_seed=42) == datasource.query(
+        limit=5, shuffle=True, shuffle_seed=42
+    )
+
+
+def test_query_datasource_with_filters():
+    """Test querying datasource rows with various filter operators."""
+    # Create a datasource with test data
+    test_data = [
+        {"name": "Alice", "age": 25, "city": "New York", "score": 85.5},
+        {"name": "Bob", "age": 30, "city": "San Francisco", "score": 90.0},
+        {"name": "Charlie", "age": 35, "city": "Chicago", "score": 75.5},
+        {"name": "Diana", "age": 28, "city": "Boston", "score": 88.0},
+        {"name": "Eve", "age": 32, "city": "New York", "score": 92.0},
+    ]
+    datasource = Datasource.from_list(name=f"test_filter_datasource_{uuid4()}", data=test_data)
+
+    # Test == operator
+    rows = datasource.query(filters=[("city", "==", "New York")])
+    assert len(rows) == 2
+    assert all(row["city"] == "New York" for row in rows)
+
+    # Test > operator
+    rows = datasource.query(filters=[("age", ">", 30)])
+    assert len(rows) == 2
+    assert all(row["age"] > 30 for row in rows)
+
+    # Test >= operator
+    rows = datasource.query(filters=[("score", ">=", 88.0)])
+    assert len(rows) == 3
+    assert all(row["score"] >= 88.0 for row in rows)
+
+    # Test < operator
+    rows = datasource.query(filters=[("age", "<", 30)])
+    assert len(rows) == 2
+    assert all(row["age"] < 30 for row in rows)
+
+    # Test in operator
+    rows = datasource.query(filters=[("city", "in", ["New York", "Boston"])])
+    assert len(rows) == 3
+    assert all(row["city"] in ["New York", "Boston"] for row in rows)
+
+    # Test not in operator
+    rows = datasource.query(filters=[("city", "not in", ["New York", "Boston"])])
+    assert len(rows) == 2
+    assert all(row["city"] not in ["New York", "Boston"] for row in rows)
+
+    # Test like operator
+    rows = datasource.query(filters=[("name", "like", "li")])
+    assert len(rows) == 2
+    assert all("li" in row["name"].lower() for row in rows)
+
+    # Test multiple filters (AND logic)
+    rows = datasource.query(filters=[("city", "==", "New York"), ("age", ">", 26)])
+    assert len(rows) == 1
+    assert rows[0]["name"] == "Eve"
+
+    # Test filter with pagination
+    rows = datasource.query(filters=[("age", ">=", 28)], limit=2, offset=1)
+    assert len(rows) == 2
+
+
+def test_query_datasource_with_none_filters():
+    """Test filtering for None values."""
+    test_data = [
+        {"name": "Alice", "age": 25, "label": "A"},
+        {"name": "Bob", "age": 30, "label": None},
+        {"name": "Charlie", "age": 35, "label": "C"},
+        {"name": "Diana", "age": None, "label": "D"},
+        {"name": "Eve", "age": 32, "label": None},
+    ]
+    datasource = Datasource.from_list(name=f"test_none_filter_{uuid4()}", data=test_data)
+
+    # Test == None
+    rows = datasource.query(filters=[("label", "==", None)])
+    assert len(rows) == 2
+    assert all(row["label"] is None for row in rows)
+
+    # Test != None
+    rows = datasource.query(filters=[("label", "!=", None)])
+    assert len(rows) == 3
+    assert all(row["label"] is not None for row in rows)
+
+    # Test that None values are excluded from comparison operators
+    rows = datasource.query(filters=[("age", ">", 25)])
+    assert len(rows) == 3
+    assert all(row["age"] is not None and row["age"] > 25 for row in rows)
+
+
+def test_query_datasource_filter_invalid_column():
+    """Test that querying with an invalid column raises an error."""
+    test_data = [{"name": "Alice", "age": 25}]
+    datasource = Datasource.from_list(name=f"test_invalid_filter_{uuid4()}", data=test_data)
+
+    with pytest.raises(ValueError):
+        datasource.query(filters=[("invalid_column", "==", "test")])
+
+
+def test_to_list(hf_dataset, datasource):
+    assert datasource.to_list() == hf_dataset.to_list()
+
+
+def test_download_datasource(hf_dataset, datasource):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Dataset download
+        datasource.download(temp_dir)
+        downloaded_hf_dataset_dir = f"{temp_dir}/{datasource.name}"
+        assert os.path.exists(downloaded_hf_dataset_dir)
+        assert os.path.isdir(downloaded_hf_dataset_dir)
+        assert not os.path.exists(f"{downloaded_hf_dataset_dir}.zip")
+        dataset_from_downloaded_hf_dataset = Dataset.load_from_disk(downloaded_hf_dataset_dir)
+        assert dataset_from_downloaded_hf_dataset.column_names == hf_dataset.column_names
+        assert dataset_from_downloaded_hf_dataset.to_dict() == hf_dataset.to_dict()
+
+        # JSON download
+        datasource.download(temp_dir, file_type="json")
+        downloaded_json_file = f"{temp_dir}/{datasource.name}.json"
+        assert os.path.exists(downloaded_json_file)
+        with open(downloaded_json_file, "r") as f:
+            content = json.load(f)
+            assert content == hf_dataset.to_list()
+
+        # CSV download
+        datasource.download(temp_dir, file_type="csv")
+        downloaded_csv_file = f"{temp_dir}/{datasource.name}.csv"
+        assert os.path.exists(downloaded_csv_file)
+        dataset_from_downloaded_csv = cast(Dataset, Dataset.from_csv(downloaded_csv_file))
+        assert dataset_from_downloaded_csv.column_names == hf_dataset.column_names
+        assert (
+            dataset_from_downloaded_csv.remove_columns("score").to_dict()
+            == hf_dataset.remove_columns("score").to_dict()
+        )
+        # Replace None with NaN for comparison
+        assert np.allclose(
+            np.array([np.nan if v is None else float(v) for v in dataset_from_downloaded_csv["score"]], dtype=float),
+            np.array([np.nan if v is None else float(v) for v in hf_dataset["score"]], dtype=float),
+            equal_nan=True,
+        )