Consolidate/structure unit tests

[Avi-Kenny]

The unit tests are sprawling and could benefit from additional structure. Also, we need to check code coverage and write tests for uncovered functionality.