Add unit test (forgot it from the previous commit)