Refactored and added driver test code