<!-- Please only use this template for submitting enhancement requests --> **What would you like to be added**: Add UT with pytest. **Why is this needed**:
