Improved previous fix (Solved small issue that makes test fail)