Reverted change that broke tests