Updated test & fixed bugs
- Couple of bugs found when fixing the tests.
- Mainly returning incorrect keys