Add two more UTs and some minor optimisation changes