Fixed bugs so it passes go test