Added tests - still not working correctly