FIX: Speed up and make test more reliable