Add missing count check in test