import pytest
import importlib
import xarray as xr
import dask
import numpy as np
code = importlib.import_module('preproc.interp')

# def test_plot_station():
#     assert code.plot_station(0) == 0
#     assert code.plot_station(0) == 1
#
def test_preprocess():
    # Create sample dataset with 10 timesteps
    data = xr.DataArray(
        np.random.randn(10, 3, 3),
        dims=("time", "lat", "lon"),
        coords={"time": range(10), "lat": [10, 20, 30], "lon": [-120, -110, -100]},
    )
    ds = xr.Dataset({"var": data})

    # Test that preprocess function keeps only first 5 timesteps
    preprocessed_ds = code.preprocess(ds, n=5)
    assert preprocessed_ds.dims == {"time": 5, "lat": 3, "lon": 3}
    assert preprocessed_ds.coords["time"].values.tolist() == [0, 1, 2, 3, 4]

    # Test that preprocess function keeps only first 3 timesteps
    preprocessed_ds = code.preprocess(ds, n=3)
    assert preprocessed_ds.dims == {"time": 3, "lat": 3, "lon": 3}
    assert preprocessed_ds.coords["time"].values.tolist() == [0, 1, 2]
