update integration tests

sdv-dev · Feb 11, 2022 · 129b19f · 129b19f
1 parent 08f85ce
commit 129b19f
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 85 deletions.
diff --git a/tests/integration/tabular/test_base.py b/tests/integration/tabular/test_base.py
@@ -38,23 +38,23 @@ def test_conditional_sampling_graceful_reject_sampling_True_dict(model):
         model.sample_conditions(conditions=conditions)
 
 
-#@pytest.mark.parametrize('model', MODELS)
-#def test_conditional_sampling_graceful_reject_sampling_True_dataframe(model):
-#    data = pd.DataFrame({
-#        'column1': list(range(100)),
-#        'column2': list(range(100)),
-#        'column3': list(range(100))
-#    })
-#
-#    model.fit(data)
-#    conditions = pd.DataFrame({
-#        'column1': [28],
-#        'column2': [37],
-#        'column3': [93]
-#    })
-#
-#    with pytest.raises(ValueError):
-#        model.sample(conditions=conditions, graceful_reject_sampling=True)
+@pytest.mark.parametrize('model', MODELS)
+def test_conditional_sampling_graceful_reject_sampling_True_dataframe(model):
+    data = pd.DataFrame({
+        'column1': list(range(100)),
+        'column2': list(range(100)),
+        'column3': list(range(100))
+    })
+
+    model.fit(data)
+    conditions = pd.DataFrame({
+        'column1': [28],
+        'column2': [37],
+        'column3': [93]
+    })
+
+    with pytest.raises(ValueError):
+        model.sample_remaining_columns(conditions)
 
 
 def test_fit_with_unique_constraint_on_data_with_only_index_column():

diff --git a/tests/integration/tabular/test_copulagan.py b/tests/integration/tabular/test_copulagan.py
@@ -107,23 +107,23 @@ def test_conditional_sampling_dict():
     assert set(sampled["column2"].unique()) == set(["b"])
 
 
-#def test_conditional_sampling_dataframe():
-#    data = pd.DataFrame({
-#        "column1": [1.0, 0.5, 2.5] * 10,
-#        "column2": ["a", "b", "c"] * 10
-#    })
-#
-#    model = CopulaGAN(epochs=1)
-#    model.fit(data)
-#    conditions = pd.DataFrame({
-#        "column2": ["b", "b", "b", "c", "c"]
-#    })
-#    sampled = model.sample(conditions=conditions)
-#
-#    assert sampled.shape[0] == len(conditions["column2"])
-#    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
-#
-#
+def test_conditional_sampling_dataframe():
+    data = pd.DataFrame({
+        "column1": [1.0, 0.5, 2.5] * 10,
+        "column2": ["a", "b", "c"] * 10
+    })
+
+    model = CopulaGAN(epochs=1)
+    model.fit(data)
+    conditions = pd.DataFrame({
+        "column2": ["b", "b", "b", "c", "c"]
+    })
+    sampled = model.sample_remaining_columns(conditions)
+
+    assert sampled.shape[0] == len(conditions["column2"])
+    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
+
+
 def test_conditional_sampling_two_conditions():
     data = pd.DataFrame({
         "column1": [1.0, 0.5, 2.5] * 10,

diff --git a/tests/integration/tabular/test_copulas.py b/tests/integration/tabular/test_copulas.py
@@ -200,23 +200,23 @@ def test_conditional_sampling_dict():
     assert set(sampled["column2"].unique()) == set(["b"])
 
 
-#def test_conditional_sampling_dataframe():
-#    data = pd.DataFrame({
-#        "column1": [1.0, 0.5, 2.5] * 10,
-#        "column2": ["a", "b", "c"] * 10
-#    })
-#
-#    model = GaussianCopula()
-#    model.fit(data)
-#    conditions = pd.DataFrame({
-#        "column2": ["b", "b", "b", "c", "c"]
-#    })
-#    sampled = model.sample(conditions=conditions)
-#
-#    assert sampled.shape[0] == len(conditions["column2"])
-#    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
-#
-#
+def test_conditional_sampling_dataframe():
+    data = pd.DataFrame({
+        "column1": [1.0, 0.5, 2.5] * 10,
+        "column2": ["a", "b", "c"] * 10
+    })
+
+    model = GaussianCopula()
+    model.fit(data)
+    conditions = pd.DataFrame({
+        "column2": ["b", "b", "b", "c", "c"]
+    })
+    sampled = model.sample_remaining_columns(conditions)
+
+    assert sampled.shape[0] == len(conditions["column2"])
+    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
+
+
 def test_conditional_sampling_two_conditions():
     data = pd.DataFrame({
         "column1": [1.0, 0.5, 2.5] * 10,

diff --git a/tests/integration/tabular/test_ctgan.py b/tests/integration/tabular/test_ctgan.py
@@ -104,23 +104,23 @@ def test_conditional_sampling_dict():
     assert set(sampled["column2"].unique()) == set(["b"])
 
 
-#def test_conditional_sampling_dataframe():
-#    data = pd.DataFrame({
-#        "column1": [1.0, 0.5, 2.5] * 10,
-#        "column2": ["a", "b", "c"] * 10
-#    })
-#
-#    model = CTGAN(epochs=1)
-#    model.fit(data)
-#    conditions = pd.DataFrame({
-#        "column2": ["b", "b", "b", "c", "c"]
-#    })
-#    sampled = model.sample(conditions=conditions)
-#
-#    assert sampled.shape[0] == len(conditions["column2"])
-#    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
-#
-#
+def test_conditional_sampling_dataframe():
+    data = pd.DataFrame({
+        "column1": [1.0, 0.5, 2.5] * 10,
+        "column2": ["a", "b", "c"] * 10
+    })
+
+    model = CTGAN(epochs=1)
+    model.fit(data)
+    conditions = pd.DataFrame({
+        "column2": ["b", "b", "b", "c", "c"]
+    })
+    sampled = model.sample_remaining_columns(conditions)
+
+    assert sampled.shape[0] == len(conditions["column2"])
+    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
+
+
 def test_conditional_sampling_two_conditions():
     data = pd.DataFrame({
         "column1": [1.0, 0.5, 2.5] * 10,

diff --git a/tests/integration/tabular/test_tvae.py b/tests/integration/tabular/test_tvae.py
@@ -103,23 +103,23 @@ def test_conditional_sampling_dict():
     assert set(sampled["column2"].unique()) == set(["b"])
 
 
-#def test_conditional_sampling_dataframe():
-#    data = pd.DataFrame({
-#        "column1": [1.0, 0.5, 2.5] * 10,
-#        "column2": ["a", "b", "c"] * 10
-#    })
-#
-#    model = TVAE(epochs=1)
-#    model.fit(data)
-#    conditions = pd.DataFrame({
-#        "column2": ["b", "b", "b", "c", "c"]
-#    })
-#    sampled = model.sample(conditions=conditions)
-#
-#    assert sampled.shape[0] == len(conditions["column2"])
-#    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
-#
-#
+def test_conditional_sampling_dataframe():
+    data = pd.DataFrame({
+        "column1": [1.0, 0.5, 2.5] * 10,
+        "column2": ["a", "b", "c"] * 10
+    })
+
+    model = TVAE(epochs=1)
+    model.fit(data)
+    conditions = pd.DataFrame({
+        "column2": ["b", "b", "b", "c", "c"]
+    })
+    sampled = model.sample_remaining_columns(conditions)
+
+    assert sampled.shape[0] == len(conditions["column2"])
+    assert (sampled["column2"] == np.array(["b", "b", "b", "c", "c"])).all()
+
+
 def test_conditional_sampling_two_conditions():
     data = pd.DataFrame({
         "column1": [1.0, 0.5, 2.5] * 10,