Orgnize test code a little bit better

Change-Id: I10358459cceb290b85770c35a798282f9ed46140