Skip to content

Commit 67f075a

Browse files
ahuang11ahuang11
authored
Migrate tests (#155)
Closes #150 Moves inside `tests/` and renames them to `test_` Also add conftest for fixtures.
1 parent f1f6694 commit 67f075a

6 files changed

Lines changed: 154 additions & 147 deletions

File tree

‎pyproject.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ features = ["pyo3/extension-module"]
5151
module-name = "xarray_sql._native"
5252

5353
[tool.setuptools.packages.find]
54-
exclude = ["demo", "perf_tests"]
54+
exclude = ["demo", "perf_tests", "tests", "tests.*"]
5555

5656
[tool.pyink]
5757
line-length = 80
@@ -93,3 +93,6 @@ dev = [
9393
[tool.uv]
9494
# Rebuild package when any rust files change
9595
cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "**/*.rs"}]
96+
97+
[tool.pytest.ini_options]
98+
testpaths = ["tests"]

‎tests/__init__.py‎

Whitespace-only changes.

‎tests/conftest.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
import pytest
2+
3+
import numpy as np
4+
import pandas as pd
5+
import xarray as xr
6+
7+
8+
def rand_wx(start: str, end: str) -> xr.Dataset:
9+
np.random.seed(42)
10+
lat = np.linspace(-90, 90, num=720)
11+
lon = np.linspace(-180, 180, num=1440)
12+
time = pd.date_range(start, end, freq="h")
13+
level = np.array([1000, 500], dtype=np.int32)
14+
reference_time = pd.Timestamp(start)
15+
temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
16+
precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
17+
return xr.Dataset(
18+
data_vars=dict(
19+
temperature=(["lat", "lon", "time", "level"], temperature),
20+
precipitation=(["lat", "lon", "time", "level"], precipitation),
21+
),
22+
coords=dict(
23+
lat=lat,
24+
lon=lon,
25+
time=time,
26+
level=level,
27+
reference_time=reference_time,
28+
),
29+
attrs=dict(description="Random weather."),
30+
)
31+
32+
33+
def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
34+
"""Create a large xarray dataset for memory testing."""
35+
np.random.seed(42)
36+
37+
time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
38+
lat = np.linspace(-90, 90, lat_points)
39+
lon = np.linspace(-180, 180, lon_points)
40+
41+
temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
42+
precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
43+
44+
return xr.Dataset(
45+
{
46+
"temperature": (["time", "lat", "lon"], temp_data),
47+
"precipitation": (["time", "lat", "lon"], precip_data),
48+
},
49+
coords={"time": time, "lat": lat, "lon": lon},
50+
)
51+
52+
53+
@pytest.fixture
54+
def air():
55+
ds = xr.tutorial.open_dataset("air_temperature")
56+
chunks = {"time": 240}
57+
return ds.chunk(chunks)
58+
59+
60+
@pytest.fixture
61+
def air_small(air):
62+
return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk(
63+
{"time": 240}
64+
)
65+
66+
67+
@pytest.fixture
68+
def randwx():
69+
return rand_wx("1995-01-13T00", "1995-01-13T01")
70+
71+
72+
@pytest.fixture
73+
def large_ds():
74+
return create_large_dataset().chunk({"time": 25})
75+
76+
77+
@pytest.fixture
78+
def air_dataset_small():
79+
ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
80+
return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
81+
82+
83+
@pytest.fixture
84+
def air_dataset_large():
85+
return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
86+
87+
88+
@pytest.fixture
89+
def weather_dataset():
90+
ds = rand_wx("2023-01-01T00", "2023-01-01T12")
91+
return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
92+
{"time": 3}
93+
)
94+
95+
96+
@pytest.fixture
97+
def synthetic_dataset():
98+
return create_large_dataset(
99+
time_steps=50, lat_points=20, lon_points=20
100+
).chunk({"time": 25})
101+
102+
103+
@pytest.fixture
104+
def station_dataset():
105+
return xr.Dataset(
106+
{
107+
"station_id": (["station"], [1, 2, 3, 4, 5]),
108+
"elevation": (["station"], [100, 250, 500, 750, 1000]),
109+
"name": (
110+
["station"],
111+
["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"],
112+
),
113+
}
114+
).chunk({"station": 5})
115+
116+
117+
@pytest.fixture
118+
def air_and_stations():
119+
air = (
120+
xr.tutorial.open_dataset("air_temperature")
121+
.isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
122+
.chunk({"time": 6})
123+
)
124+
stations = xr.Dataset(
125+
{
126+
"station_id": (["station"], [101, 102, 103]),
127+
"lat": (
128+
["station"],
129+
[air.lat.values[0], air.lat.values[2], air.lat.values[4]],
130+
),
131+
"lon": (
132+
["station"],
133+
[air.lon.values[1], air.lon.values[3], air.lon.values[5]],
134+
),
135+
"elevation": (["station"], [100, 250, 500]),
136+
}
137+
).chunk({"station": 3})
138+
return air, stations

xarray_sql/df_test.py renamed to tests/test_df.py

Lines changed: 8 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import pytest
77
import xarray as xr
88

9-
from .df import (
9+
from xarray_sql.df import (
1010
DEFAULT_BATCH_SIZE,
1111
_parse_schema,
1212
block_slices,
@@ -17,82 +17,7 @@
1717
iter_record_batches,
1818
pivot,
1919
)
20-
from .reader import read_xarray, read_xarray_table
21-
22-
23-
def rand_wx(start: str, end: str) -> xr.Dataset:
24-
np.random.seed(42)
25-
lat = np.linspace(-90, 90, num=720)
26-
lon = np.linspace(-180, 180, num=1440)
27-
time = pd.date_range(start, end, freq="h")
28-
level = np.array([1000, 500], dtype=np.int32)
29-
reference_time = pd.Timestamp(start)
30-
temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level))
31-
precipitation = 10 * np.random.rand(720, 1440, len(time), len(level))
32-
return xr.Dataset(
33-
data_vars=dict(
34-
temperature=(["lat", "lon", "time", "level"], temperature),
35-
precipitation=(["lat", "lon", "time", "level"], precipitation),
36-
),
37-
coords=dict(
38-
lat=lat,
39-
lon=lon,
40-
time=time,
41-
level=level,
42-
reference_time=reference_time,
43-
),
44-
attrs=dict(description="Random weather."),
45-
)
46-
47-
48-
def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100):
49-
"""Create a large xarray dataset for memory testing."""
50-
np.random.seed(42)
51-
52-
time = pd.date_range("2020-01-01", periods=time_steps, freq="h")
53-
lat = np.linspace(-90, 90, lat_points)
54-
lon = np.linspace(-180, 180, lon_points)
55-
56-
temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10
57-
precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100
58-
59-
return xr.Dataset(
60-
{
61-
"temperature": (["time", "lat", "lon"], temp_data),
62-
"precipitation": (["time", "lat", "lon"], precip_data),
63-
},
64-
coords={"time": time, "lat": lat, "lon": lon},
65-
)
66-
67-
68-
def adding_function(x, y):
69-
"""Simple function that adds two values and returns a DataFrame."""
70-
result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]})
71-
return result
72-
73-
74-
@pytest.fixture
75-
def air():
76-
ds = xr.tutorial.open_dataset("air_temperature")
77-
chunks = {"time": 240}
78-
return ds.chunk(chunks)
79-
80-
81-
@pytest.fixture
82-
def air_small(air):
83-
return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk(
84-
{"time": 240}
85-
)
86-
87-
88-
@pytest.fixture
89-
def randwx():
90-
return rand_wx("1995-01-13T00", "1995-01-13T01")
91-
92-
93-
@pytest.fixture
94-
def large_ds():
95-
return create_large_dataset().chunk({"time": 25})
20+
from xarray_sql.reader import read_xarray, read_xarray_table
9621

9722

9823
def test_explode_cardinality(air):
@@ -295,6 +220,12 @@ def test_from_map_batched_basic_functionality(air_small):
295220
assert len(batch) > 0
296221

297222

223+
def adding_function(x, y):
224+
"""Simple function that adds two values and returns a DataFrame."""
225+
result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]})
226+
return result
227+
228+
298229
def test_from_map_batched_multiple_iterables():
299230
x_values = [1, 2, 3, 4, 5]
300231
y_values = [10, 20, 30, 40, 50]

xarray_sql/reader_test.py renamed to tests/test_reader.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,9 @@
2323
import xarray as xr
2424
from datafusion import SessionContext
2525

26-
from ._native import LazyArrowStreamTable
27-
from .reader import XarrayRecordBatchReader, read_xarray_table
28-
from .df import _parse_schema
26+
from xarray_sql._native import LazyArrowStreamTable
27+
from xarray_sql.reader import XarrayRecordBatchReader, read_xarray_table
28+
from xarray_sql.df import _parse_schema
2929

3030

3131
@pytest.fixture

xarray_sql/sql_test.py renamed to tests/test_sql.py

Lines changed: 1 addition & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -3,72 +3,7 @@
33
import pytest
44
import xarray as xr
55

6-
from . import XarrayContext
7-
from .df_test import create_large_dataset, rand_wx
8-
9-
10-
@pytest.fixture
11-
def air_dataset_small():
12-
ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
13-
return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10))
14-
15-
16-
@pytest.fixture
17-
def air_dataset_large():
18-
return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240})
19-
20-
21-
@pytest.fixture
22-
def weather_dataset():
23-
ds = rand_wx("2023-01-01T00", "2023-01-01T12")
24-
return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk(
25-
{"time": 3}
26-
)
27-
28-
29-
@pytest.fixture
30-
def synthetic_dataset():
31-
return create_large_dataset(
32-
time_steps=50, lat_points=20, lon_points=20
33-
).chunk({"time": 25})
34-
35-
36-
@pytest.fixture
37-
def station_dataset():
38-
return xr.Dataset(
39-
{
40-
"station_id": (["station"], [1, 2, 3, 4, 5]),
41-
"elevation": (["station"], [100, 250, 500, 750, 1000]),
42-
"name": (
43-
["station"],
44-
["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"],
45-
),
46-
}
47-
).chunk({"station": 5})
48-
49-
50-
@pytest.fixture
51-
def air_and_stations():
52-
air = (
53-
xr.tutorial.open_dataset("air_temperature")
54-
.isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8))
55-
.chunk({"time": 6})
56-
)
57-
stations = xr.Dataset(
58-
{
59-
"station_id": (["station"], [101, 102, 103]),
60-
"lat": (
61-
["station"],
62-
[air.lat.values[0], air.lat.values[2], air.lat.values[4]],
63-
),
64-
"lon": (
65-
["station"],
66-
[air.lon.values[1], air.lon.values[3], air.lon.values[5]],
67-
),
68-
"elevation": (["station"], [100, 250, 500]),
69-
}
70-
).chunk({"station": 3})
71-
return air, stations
6+
from xarray_sql import XarrayContext
727

738

749
def test_sanity(air_dataset_small):

0 commit comments

Comments
 (0)