Updating source

Removed test code from main class
Added a test class