Added tests to flatten