Rewrote a test
Add integration test