@mohit I added a bunch more tests (and fixes) over several commits and I think I've addressed this. Can you re-review?