Cleanup and unit test
Fixed NPE