[Fix] Fix unit tests, all green again