'Verifies that the given unit\'s messages match the expected messages.
If expected_status_message is None, then the unit status message is expected to not be present.'

def _verify_unit_warning(self, unit_state, expected_status_message):

    self._ensure_unit_present(unit_state)
    self.course_outline_page.visit()
    section = self.course_outline_page.section(unit_state.name)
    subsection = section.subsection_at(0)
    subsection.expand_subsection()
    unit = subsection.unit_at(0)
    if (expected_status_message == self.STAFF_ONLY_WARNING):
        self.assertEqual(section.status_message, self.STAFF_ONLY_WARNING)
        self.assertEqual(subsection.status_message, self.STAFF_ONLY_WARNING)
        self.assertEqual(unit.status_message, self.STAFF_ONLY_WARNING)
    else:
        self.assertFalse(section.has_status_message)
        self.assertFalse(subsection.has_status_message)
        if expected_status_message:
            self.assertEqual(unit.status_message, expected_status_message)
        else:
            self.assertFalse(unit.has_status_message)
