Merge "Refactor function signature and move mock back" into main