Update for test cleanup