Add new test
Fix failing test
Add integration tests