Couple of bugs from refactor
2 files changed