Re-added test fix that was accidentally removed