Add test with chunking

Author: eric-czech

sgkit/tests/test_association.py

@@ -9,7 +9,6 @@
 from pandas import DataFrame
 from xarray import Dataset
 
-from sgkit import variables
 from sgkit.stats.association import gwas_linear_regression, linear_regression
 from sgkit.typing import ArrayLike
 
@@ -175,16 +174,24 @@ def validate(dfp: DataFrame, dft: DataFrame) -> None:
 
 def test_gwas_linear_regression__lazy_results(ds):
     res = gwas_linear_regression(
-        ds, dosage="dosage", covariates="covar_0", traits="trait_0"
+        ds, dosage="dosage", covariates="covar_0", traits="trait_0", merge=False
     )
-    for v in [
-        variables.variant_beta,
-        variables.variant_t_value,
-        variables.variant_p_value,
-    ]:
+    for v in res:
         assert isinstance(res[v].data, da.Array)
 
 
+@pytest.mark.parametrize("chunks", [5, -1, "auto"])
+def test_gwas_linear_regression__variable_shapes(ds, chunks):
+    ds = ds.chunk(chunks=chunks)
+    res = gwas_linear_regression(
+        ds, dosage="dosage", covariates="covar_0", traits="trait_0", merge=False
+    )
+    shape = (ds.dims["variants"], 1)
+    for v in res:
+        assert res[v].data.shape == shape
+        assert res[v].data.compute().shape == shape
+
+
 def test_gwas_linear_regression__multi_trait(ds):
     def run(traits: Sequence[str]) -> Dataset:
         return gwas_linear_regression(