Updated tests as part of #304