Some code refactoring to make automation api easier to use