Add some tests

[Avinash]

I was going through the code base and didn't saw any test. I can start with adding some unit tests and integration test as well.