Code cleanup
Fixed test
Minor change