Make cli module more flexible
Add tests for cli module