- fixed the failing tests
- more code cleanup