def test__get_extensions(self):
"""Test get list of extensions from childs"""
# Setup
modeler = Mock()
model_dict = [{
'model': 'data 1'
}, {
'model': 'data 2'
}, {
'model': 'data 3'
}]
modeler._get_model_dict.side_effect = model_dict
# Run
child_table = pd.DataFrame({'foo': ['aaa', 'bbb', 'ccc']})
result = Modeler._get_extension(modeler, 'some_name', child_table,
'foo')
# Asserts
expected = pd.DataFrame(
{
'__some_name__model': ['data 1', 'data 2', 'data 3'],
'__some_name__child_rows': [1, 1, 1]
},
index=['aaa', 'bbb', 'ccc'])
pd.testing.assert_frame_equal(result, expected)
assert modeler._get_model_dict.call_count == 3
def test__get_extensions(self):
"""Test get list of extensions from childs"""
# Setup
model = Mock(spec=SDVModel)
model.return_value = model
model.get_parameters.side_effect = [{
'model': 'data 1'
}, {
'model': 'data 2'
}, {
'model': 'data 3'
}]
modeler = Mock(spec=Modeler)
modeler.model = model
modeler.model_kwargs = dict()
modeler.metadata = Mock(spec=Metadata)
# Run
child_table = pd.DataFrame({'foo': ['aaa', 'bbb', 'ccc']})
result = Modeler._get_extension(modeler, 'some_name', child_table,
'foo')
# Asserts
expected = pd.DataFrame(
{
'__some_name__model': ['data 1', 'data 2', 'data 3'],
'__some_name__child_rows': [1, 1, 1]
},
index=['aaa', 'bbb', 'ccc'])
pd.testing.assert_frame_equal(result, expected)
assert model.get_parameters.call_count == 3
