Fixes for broken test