Added some more unit tests + cleanup fixes