Fixes for issues caught by tests

@ransomr Can you please review.

Thanks!
Heman