Cleaned up the internal logic of the bandit testing plugin