class_docstr: Test cases for diagnostics infra.
class_docstr_tok: ["Test","cases","for","diagnostics","infra","."]
code: def test_diagnostics_engine_records_diagnosis_with_custom_rules(self):
      custom_rules = infra.RuleCollection.custom_collection_from_list(
          "CustomRuleCollection",
          [
              infra.Rule(
                  "1",
                  "custom-rule",
                  message_default_template="custom rule message",
              ),
              infra.Rule(
                  "2",
                  "custom-rule-2",
                  message_default_template="custom rule message 2",
              ),
          ],
      )
  
      with assert_all_diagnostics(
          self,
          self.context,
          {
              (custom_rules.custom_rule, infra.Level.WARNING),  # type: ignore[attr-defined]
              (custom_rules.custom_rule_2, infra.Level.ERROR),  # type: ignore[attr-defined]
          },
      ):
          diagnostic1 = infra.Diagnostic(
              custom_rules.custom_rule, infra.Level.WARNING  # type: ignore[attr-defined]
          )
          self.context.log(diagnostic1)
  
          diagnostic2 = infra.Diagnostic(
              custom_rules.custom_rule_2, infra.Level.ERROR  # type: ignore[attr-defined]
          )
          self.context.log(diagnostic2)
  
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 the TestDiagnosticsInfra obj's underlying member method code with example usage: obj.test_diagnostics_engine_records_diagnosis_with_custom_rules() without return types