code: def context(model_dir, model_name):
      micro_batching_params = {
          "mb_size": 2,
          "mb_parallelism": {
              "preprocess": 1,
              "inference": 2,
              "postprocess": 3,
          },
      }
  
      config_file = Path(model_dir).joinpath("micro_batching.json")
  
      with open(config_file, "w") as f:
          json.dump(micro_batching_params, f)
  
      context = MockContext(
          model_name="mnist",
          model_dir=model_dir.as_posix(),
          model_file=model_name + ".py",
      )
      context.model_yaml_config = micro_batching_params
      yield context
  
imports: import json
  import math
  import random
  import sys
  from pathlib import Path
  import pytest
  from torchvision.models.resnet import ResNet18_Weights
  from ts.handler_utils.micro_batching import MicroBatching
  from ts.torch_handler.image_classifier import ImageClassifier
  from ts.torch_handler.unit_tests.test_utils.mock_context import MockContext
  from ts.torch_handler.unit_tests.test_utils.model_dir import copy_files, download_model
label_desc: Use image node_id 6 for calling a global function with example usage: context(model_dir, model_name) without return types