def test_sample_remaining_columns(self):
"""Test ``sample_remaining_columns`` method.
Expect the correct args to be passed to ``_sample_remaining_columns``
Input:
- valid DataFrame
Side Effects:
- The expected ``_sample_remaining_columns`` call.
"""
# Setup
model = Mock(spec_set=GaussianCopula)
conditions = pd.DataFrame([{'cola': 'a'}] * 5)
batch_size = 1
randomize_samples = False
output_file_path = 'test.csv'
# Run
out = GaussianCopula.sample_remaining_columns(
model,
conditions,
batch_size=batch_size,
randomize_samples=randomize_samples,
output_file_path=output_file_path,
)
# Assert
model._sample_remaining_columns.assert_called_once_with(
conditions, 100, batch_size, randomize_samples, output_file_path)
assert out == model._sample_remaining_columns.return_value
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()