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