Fix tests
I'm an idiot