More fixing of test code