Add extra unit test for COUNT and TotalSegments