refact code for tests