Minor bug fix and added a test case!