fix tests. not beautiful but does the job for now.