[Refactor] - KD tree clean up and unit tests

[danielduhh]

We need unit tests for our KD tree class. This will also requires abstracting out some of the logic to make it more general