Refactor test to support more test variation

Change-Id: I61c8626bc78ae6ffeb10665a82a89881985c80c8