Test case improvements and minor code cleanup