Changed some testing code in main()