rewrite tests a bit