Added more tests. Code cleanup.