Added basic unit testing