def METHOD_NAME(self):
    try:
        tf2.summary.experimental.set_step(333)
        # TODO(nickfelt): change test logic so we can just omit `step` entirely.
        event = self.scalar_event("a", 1.0, step=None)
        self.assertEqual(333, event.step)
    finally:
        # Reset to default state for other tests.
        tf2.summary.experimental.set_step(None)