def METHOD_NAME(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)
    METHOD_NAME = MockContext(
        model_name="mnist",
        model_dir=model_dir.as_posix(),
        model_file=model_name + ".py",
    )
    METHOD_NAME.model_yaml_config = micro_batching_params
    yield METHOD_NAME