Rewrite tests a bit for more clarity and reuse
Cover websockets