Unit test

Hi Rob, 

I have added a couple of unit tests and the previous review changes. 

Shall we push these in the interim? 

Thanks
Sagnik