Add back some more tests