[DI] Update unit tests