Refactor & more hacky code to test