Fixes broken test