address comments: add test