Made some small refactorings

Tests are still passing after a small refactoring to remove a line from being repeated. I couldn't see the need for any additional tests to be added but if you need some tests please shout