Update unit and integration test