Refactor code for using class
Add class wrapper to get multiple context