Re-adding unit test for corrected code

Glad I did the test. Next time I'll have to set up the unit test environment locally so there isn't so much pull request spam.