Code cleanup
Fixed tests