Fix tests for flattening.