Add unit tests
Add integration tests