Tests are very nice!
I'd be tempted to merge it already, but will let you decide if you want to polish on the other comments (there are no blockers)