Fix up test case and add another