Fixed test class so it actually worked