remove extra code from main function and add into doctest
remove extra code from main function and add into doctest