code: def assert_all_diagnostics(
      test_suite: unittest.TestCase,
      sarif_log_builder: _SarifLogBuilder,
      rule_level_pairs: AbstractSet[Tuple[infra.Rule, infra.Level]],
  ):
      """Context manager to assert that all diagnostics are emitted.
  
      Usage:
          with assert_all_diagnostics(
              self,
              diagnostics.engine,
              {(rule, infra.Level.Error)},
          ):
              torch.onnx.export(...)
  
      Args:
          test_suite: The test suite instance.
          sarif_log_builder: The SARIF log builder.
          rule_level_pairs: A set of rule and level pairs to assert.
  
      Returns:
          A context manager.
  
      Raises:
          AssertionError: If not all diagnostics are emitted.
      """
  
      try:
          yield
      except errors.OnnxExporterError:
          test_suite.assertIn(
              infra.Level.ERROR,
              {level for _, level in rule_level_pairs},
          )
      finally:
          _assert_has_diagnostics(sarif_log_builder, rule_level_pairs)
  
imports: from __future__ import annotations
  import contextlib
  import dataclasses
  import io
  import logging
  import typing
  import unittest
  from typing import AbstractSet, Protocol, Tuple
  import torch
  from torch.onnx import errors
  from torch.onnx._internal import diagnostics
  from torch.onnx._internal.diagnostics import infra
  from torch.onnx._internal.diagnostics.infra import formatter, sarif
  from torch.onnx._internal.fx import diagnostics
  from torch.testing._internal import common_utils, logging_utils
label_desc: Use image node_id 2 for calling a global function with example usage: assert_all_diagnostics(test_suite, sarif_log_builder, rule_level_pairs) without return types