Added --test flag. Refactored the code.