mini refactor and add diagonal test