Fixing test cases. Not working properly yet