Refactor test input layer