Fix test failures

Merged in the changes from @viveksjain and fixed any test failures that came up.
I'd rather stretch this over more lines and have it look like the other if/elseif/else in the code.