- Added more tests
- Code cleanup