From 64378af2869b944713b4154e5657d1caa57835fe Mon Sep 17 00:00:00 2001
From: Felipe Alex Hofmann <fealho@gmail.com>
Date: Wed, 11 Sep 2024 09:09:44 -0700
Subject: [PATCH] Add test for `load_demo` method (#399)

---
 tests/integration/test_load_demo.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 tests/integration/test_load_demo.py

diff --git a/tests/integration/test_load_demo.py b/tests/integration/test_load_demo.py
new file mode 100644
index 00000000..e37eca7d
--- /dev/null
+++ b/tests/integration/test_load_demo.py
@@ -0,0 +1,28 @@
+from ctgan import CTGAN, load_demo
+
+
+def test_load_demo():
+    """End-to-end test to load and synthesize data."""
+    # Setup
+    discrete_columns = [
+        'workclass',
+        'education',
+        'marital-status',
+        'occupation',
+        'relationship',
+        'race',
+        'sex',
+        'native-country',
+        'income',
+    ]
+    ctgan = CTGAN(epochs=1)
+
+    # Run
+    data = load_demo()
+    ctgan.fit(data, discrete_columns)
+    samples = ctgan.sample(1000, condition_column='native-country', condition_value='United-States')
+
+    # Assert
+    assert samples.shape == (1000, 15)
+    assert all([col[0] != ' ' for col in samples.columns])
+    assert not samples.isna().any().any()