Fix undefined method errors in test cases