Unit Testing Fix