Minor code reorganization
Add test code and fix bugs