Fix up logic a bit; pass unit tests.