code: def test_on_before_zero_grad_called(tmpdir, max_steps):
      class CurrentTestModel(BoringModel):
          on_before_zero_grad_called = 0
  
          def on_before_zero_grad(self, optimizer):
              self.on_before_zero_grad_called += 1
  
      model = CurrentTestModel()
  
      trainer = Trainer(
          default_root_dir=tmpdir, max_steps=max_steps, max_epochs=2
      )
      assert model.on_before_zero_grad_called == 0
      trainer.fit(model)
      assert max_steps == model.on_before_zero_grad_called
  
      model.on_before_zero_grad_called = 0
      trainer.test(model)
      assert model.on_before_zero_grad_called == 0
  
imports: from functools import partial, update_wrapper
  from inspect import getmembers, isfunction
  from unittest import mock
  from unittest.mock import ANY, PropertyMock
  import pytest
  import torch
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.pytorch import Callback, LightningDataModule, LightningModule, Trainer, __version__
  from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel, RandomDataset
  from lightning.pytorch.utilities.model_helpers import is_overridden
  from torch import Tensor
  from torch.utils.data import DataLoader
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 1 for calling a global function with example usage: test_on_before_zero_grad_called(tmpdir, max_steps) without return types