Fix a bug that Justin caught on the demo day, and add more unit test coverage to catch similar regressions