code: def test_main(self):
      callback = imaging.FromState("test")
  
      self.assertTrue(
          callback.on_batch({torchbearer.EPOCH: 0, "test": 1}) == 1
      )
      self.assertTrue(
          callback.on_batch({torchbearer.EPOCH: 1, "testing": 1})
          is None
      )
  
imports: from unittest import TestCase
  from mock import MagicMock, patch, ANY
  import torch
  import torchbearer
  import torchbearer.callbacks.imaging
  import matplotlib.pyplot
label_desc: Use image node_id 1 for calling the TestFromState obj's underlying member method code with example usage: obj.test_main() without return types