[test] add simple unit test for the regression dataset builder

diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644 (file)
index 0000000..e69de29

diff --git a/test/test_dataset.py b/test/test_dataset.py
new file mode 100644 (file)
index 0000000..9c9bb3b
--- /dev/null
+++ b/test/test_dataset.py
@@ -0,0 +1,11 @@
+from complexity_regularizer.dataset_builder import create_regression_data
+
+def test_regression_dataset_shapes():
+    n_samples: int = 4
+    n_latents: int = 3
+    noise_factor: float = 0.0
+    model_degree = n_latents
+    dataset = create_regression_data(n_samples, n_latents, noise_factor, model_degree)
+    X, y = dataset[:]
+    assert X.shape[0] == n_samples
+    assert X.shape[1] == model_degree + 1