Clean up of code

* Removed tests as they have another branch

* Removed redundant imports