(almost) fixed tests
Fixed tests a bit