Add upgrade method
Move common code in helpers
Fix path