[FOLLOWUP][BUGFIX] Fix failing unit tests