Updates to test code