Moved test code into a function