Some improvements to test and upgrade step
Removed some old comments
Updates in tests