Added unit test to this shit