Fix memory leak, add test, slightly refactor main()