Add test for remove and find