Fixes bug with display message logic due to deprecating GenericSource. Also added more test cases to cover a bunch more edge cases