Migrate tests

pyproject.toml

@@ -51,7 +51,7 @@ features = ["pyo3/extension-module"]
 module-name = "xarray_sql._native"
 
 [tool.setuptools.packages.find]
-exclude = ["demo", "perf_tests"]
+exclude = ["demo", "perf_tests", "tests", "tests.*"]
 
 [tool.pyink]
 line-length = 80
@@ -93,3 +93,6 @@ dev = [
 [tool.uv]
 # Rebuild package when any rust files change
 cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "**/*.rs"}]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

tests/conftest.py

@@ -0,0 +1,138 @@
+import pytest
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+def rand_wx(start: str, end: str) -> xr.Dataset:
+  np.random.seed(42)
+  lat = np.linspace(-90, 90, num=720)
+  lon = np.linspace(-180, 180, num=1440)
+  time = pd.date_range(start, end, freq="h")
+  level = np.array([1000, 500], dtype=np.int32)
+  reference_time = pd.Timestamp(start)
+  temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
+  precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
+  return xr.Dataset(
+      data_vars=dict(
+          temperature=(["lat", "lon", "time", "level"], temperature),
+          precipitation=(["lat", "lon", "time", "level"], precipitation),
+      ),
+      coords=dict(
+          lat=lat,
+          lon=lon,
+          time=time,
+          level=level,
+          reference_time=reference_time,
+      ),
+      attrs=dict(description="Random weather."),
+  )
+
+
+def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
+  """Create a large xarray dataset for memory testing."""
+  np.random.seed(42)
+
+  time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
+  lat = np.linspace(-90, 90, lat_points)
+  lon = np.linspace(-180, 180, lon_points)
+
+  temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
+  precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
+
+  return xr.Dataset(
+      {
+          "temperature": (["time", "lat", "lon"], temp_data),
+          "precipitation": (["time", "lat", "lon"], precip_data),
+      },
+      coords={"time": time, "lat": lat, "lon": lon},
+  )
+
+
+@pytest.fixture
+def air():
+  ds = xr.tutorial.open_dataset("air_temperature")
+  chunks = {"time": 240}
+  return ds.chunk(chunks)
+
+
+@pytest.fixture
+def air_small(air):
+  return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk(
+      {"time": 240}
+  )
+
+
+@pytest.fixture
+def randwx():
+  return rand_wx("1995-01-13T00", "1995-01-13T01")
+
+
+@pytest.fixture
+def large_ds():
+  return create_large_dataset().chunk({"time": 25})
+
+
+@pytest.fixture
+def air_dataset_small():
+  ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
+  return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
+
+
+@pytest.fixture
+def air_dataset_large():
+  return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
+
+
+@pytest.fixture
+def weather_dataset():
+  ds = rand_wx("2023-01-01T00", "2023-01-01T12")
+  return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
+      {"time": 3}
+  )
+
+
+@pytest.fixture
+def synthetic_dataset():
+  return create_large_dataset(
+      time_steps=50, lat_points=20, lon_points=20
+  ).chunk({"time": 25})
+
+
+@pytest.fixture
+def station_dataset():
+  return xr.Dataset(
+      {
+          "station_id": (["station"], [1, 2, 3, 4, 5]),
+          "elevation": (["station"], [100, 250, 500, 750, 1000]),
+          "name": (
+              ["station"],
+              ["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"],
+          ),
+      }
+  ).chunk({"station": 5})
+
+
+@pytest.fixture
+def air_and_stations():
+  air = (
+      xr.tutorial.open_dataset("air_temperature")
+      .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
+      .chunk({"time": 6})
+  )
+  stations = xr.Dataset(
+      {
+          "station_id": (["station"], [101, 102, 103]),
+          "lat": (
+              ["station"],
+              [air.lat.values[0], air.lat.values[2], air.lat.values[4]],
+          ),
+          "lon": (
+              ["station"],
+              [air.lon.values[1], air.lon.values[3], air.lon.values[5]],
+          ),
+          "elevation": (["station"], [100, 250, 500]),
+      }
+  ).chunk({"station": 3})
+  return air, stations

xarray_sql/df_test.py → tests/test_df.py

@@ -6,7 +6,7 @@
 import pytest
 import xarray as xr
 
-from .df import (
+from xarray_sql.df import (
     DEFAULT_BATCH_SIZE,
     _parse_schema,
     block_slices,
@@ -17,82 +17,7 @@
     iter_record_batches,
     pivot,
 )
-from .reader import read_xarray, read_xarray_table
-
-
-def rand_wx(start: str, end: str) -> xr.Dataset:
-  np.random.seed(42)
-  lat = np.linspace(-90, 90, num=720)
-  lon = np.linspace(-180, 180, num=1440)
-  time = pd.date_range(start, end, freq="h")
-  level = np.array([1000, 500], dtype=np.int32)
-  reference_time = pd.Timestamp(start)
-  temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
-  precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
-  return xr.Dataset(
-      data_vars=dict(
-          temperature=(["lat", "lon", "time", "level"], temperature),
-          precipitation=(["lat", "lon", "time", "level"], precipitation),
-      ),
-      coords=dict(
-          lat=lat,
-          lon=lon,
-          time=time,
-          level=level,
-          reference_time=reference_time,
-      ),
-      attrs=dict(description="Random weather."),
-  )
-
-
-def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
-  """Create a large xarray dataset for memory testing."""
-  np.random.seed(42)
-
-  time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
-  lat = np.linspace(-90, 90, lat_points)
-  lon = np.linspace(-180, 180, lon_points)
-
-  temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
-  precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
-
-  return xr.Dataset(
-      {
-          "temperature": (["time", "lat", "lon"], temp_data),
-          "precipitation": (["time", "lat", "lon"], precip_data),
-      },
-      coords={"time": time, "lat": lat, "lon": lon},
-  )
-
-
-def adding_function(x, y):
-  """Simple function that adds two values and returns a DataFrame."""
-  result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]})
-  return result
-
-
-@pytest.fixture
-def air():
-  ds = xr.tutorial.open_dataset("air_temperature")
-  chunks = {"time": 240}
-  return ds.chunk(chunks)
-
-
-@pytest.fixture
-def air_small(air):
-  return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk(
-      {"time": 240}
-  )
-
-
-@pytest.fixture
-def randwx():
-  return rand_wx("1995-01-13T00", "1995-01-13T01")
-
-
-@pytest.fixture
-def large_ds():
-  return create_large_dataset().chunk({"time": 25})
+from xarray_sql.reader import read_xarray, read_xarray_table
 
 
 def test_explode_cardinality(air):
@@ -295,6 +220,12 @@ def test_from_map_batched_basic_functionality(air_small):
     assert len(batch) > 0
 
 
+def adding_function(x, y):
+  """Simple function that adds two values and returns a DataFrame."""
+  result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]})
+  return result
+
+
 def test_from_map_batched_multiple_iterables():
   x_values = [1, 2, 3, 4, 5]
   y_values = [10, 20, 30, 40, 50]

xarray_sql/reader_test.py → tests/test_reader.py

@@ -23,9 +23,9 @@
 import xarray as xr
 from datafusion import SessionContext
 
-from ._native import LazyArrowStreamTable
-from .reader import XarrayRecordBatchReader, read_xarray_table
-from .df import _parse_schema
+from xarray_sql._native import LazyArrowStreamTable
+from xarray_sql.reader import XarrayRecordBatchReader, read_xarray_table
+from xarray_sql.df import _parse_schema
 
 
 @pytest.fixture

xarray_sql/sql_test.py → tests/test_sql.py

@@ -3,72 +3,7 @@
 import pytest
 import xarray as xr
 
-from . import XarrayContext
-from .df_test import create_large_dataset, rand_wx
-
-
-@pytest.fixture
-def air_dataset_small():
-  ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
-  return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
-
-
-@pytest.fixture
-def air_dataset_large():
-  return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
-
-
-@pytest.fixture
-def weather_dataset():
-  ds = rand_wx("2023-01-01T00", "2023-01-01T12")
-  return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
-      {"time": 3}
-  )
-
-
-@pytest.fixture
-def synthetic_dataset():
-  return create_large_dataset(
-      time_steps=50, lat_points=20, lon_points=20
-  ).chunk({"time": 25})
-
-
-@pytest.fixture
-def station_dataset():
-  return xr.Dataset(
-      {
-          "station_id": (["station"], [1, 2, 3, 4, 5]),
-          "elevation": (["station"], [100, 250, 500, 750, 1000]),
-          "name": (
-              ["station"],
-              ["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"],
-          ),
-      }
-  ).chunk({"station": 5})
-
-
-@pytest.fixture
-def air_and_stations():
-  air = (
-      xr.tutorial.open_dataset("air_temperature")
-      .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
-      .chunk({"time": 6})
-  )
-  stations = xr.Dataset(
-      {
-          "station_id": (["station"], [101, 102, 103]),
-          "lat": (
-              ["station"],
-              [air.lat.values[0], air.lat.values[2], air.lat.values[4]],
-          ),
-          "lon": (
-              ["station"],
-              [air.lon.values[1], air.lon.values[3], air.lon.values[5]],
-          ),
-          "elevation": (["station"], [100, 250, 500]),
-      }
-  ).chunk({"station": 3})
-  return air, stations
+from xarray_sql import XarrayContext
 
 
 def test_sanity(air_dataset_small):