code: def test_spec_augment_apply_masks_to_axis(inputs):
      """
      Tests the method _apply_masks_to_axis to see if shape is kept and
      exceptions are caught
      """
  
      data_format, axis, mask_param, n_masks = inputs
      batch_src, input_shape = get_spectrogram(data_format)
  
      spec_augment = SpecAugment(
          input_shape=input_shape,
          freq_mask_param=5,
          time_mask_param=10,
          n_freq_masks=4,
          n_time_masks=3,
          mask_value=0.0,
          data_format=data_format,
      )
  
      # We force axis that will trigger NotImplementedError
      if axis not in [0, 1, 2]:
          # Check axis error
          with pytest.raises(NotImplementedError):
              # We use batch_src instead of batch_src[0] to simulate a 4D spectrogram
              inputs = (batch_src, axis, mask_param, n_masks)
              spec_augment._apply_masks_to_axis(*inputs)
  
      # We force mask_params that will trigger the ValueError. If it is not triggered, then
      # inputs are ok, so we must only test if the shapes are kept during transformation
      elif mask_param != 5:
          # Check mask_param error
          with pytest.raises(ValueError):
              inputs = (batch_src[0], axis, mask_param, n_masks)
              spec_augment._apply_masks_to_axis(*inputs)
      else:
          # Check that transformation keeps shape
          inputs = (batch_src[0], axis, mask_param, n_masks)
          mask = spec_augment._apply_masks_to_axis(*inputs)
          np.testing.assert_equal(mask.shape[axis], input_shape[axis])
  
file_path: kapre/tests/test_augmentation.py
imports: import pytest
  import numpy
  import tensorflow
  from kapre.augmentation import ChannelSwap, SpecAugment
  from kapre.backend import _CH_FIRST_STR, _CH_LAST_STR, _CH_DEFAULT_STR
  from utils import get_audio, get_spectrogram, save_load_compare
label_desc: Use image node_id 2 for calling a global function with example usage: test_spec_augment_apply_masks_to_axis(inputs) without return types