Code cleanups, some tests, now works with loop