Add unit tests before optimizations

Add unit tests before optimizations to check if the code still generate the correct result.