Fix up test code and add more test code post re-factor
