rearchitecting, and adding some unit tests