Cleanup of tests and bringing tests and new code back in line

Signed-off-by: shalott <shalott@gmail.com>