Address review comment
Add a test