Updated tests to get working