Add more tests

We need to improve the test coverage