Repository: Project-MONAI/MONAI
File: tests/test_mapping_filed.py (Python)

Code before:
        input_file = create_input_file(self.temp_dir, name)
        output_file = os.path.join(self.output_dir, name, name + "_seg.nii.gz")
        data = [{"image": input_file}]

        test_transforms = Compose([LoadImaged(keys=["image"]), EnsureChannelFirstd(keys=["image"])])

        post_transforms = Compose(
            [
                SaveImaged(
                    keys=save_keys,
                    meta_keys="image_meta_dict",
                    output_dir=self.output_dir,
                    output_postfix="seg",
                    savepath_in_metadict=True,
                ),
                WriteFileMappingd(keys=write_keys, mapping_file_path=self.mapping_file_path),
            ]
        )

        dataset = Dataset(data=data, transform=test_transforms)
        dataloader = DataLoader(dataset, batch_size=1)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = UNet(spatial_dims=3, in_channels=1, out_channels=2, channels=(16, 32), strides=(2,)).to(device)
        model.eval()

        try:
            with torch.no_grad():
                for batch_data in dataloader:
                    test_inputs = batch_data["image"].to(device)
                    roi_size = (64, 64, 64)
                    sw_batch_size = 2
                    batch_data["seg"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, model)
                    batch_data = [post_transforms(i) for i in decollate_batch(batch_data)]

            self.assertTrue(os.path.exists(self.mapping_file_path))

            with open(self.mapping_file_path) as f:
                mapping_data = json.load(f)

            self.assertEqual(len(mapping_data), len(write_keys))
            for entry in mapping_data:
                self.assertEqual(entry["input"], input_file)
                self.assertEqual(entry["output"], output_file)

        except RuntimeError as cm:
            cause_exception = cm.__cause__
            self.assertIsInstance(cause_exception, KeyError)
            self.assertIn(
                "Missing 'saved_to' key in metadata. Check SaveImage argument 'savepath_in_metadict' is True.",
                str(cause_exception),
            )

Review comment (suggestion):
Instead of using try-except in test case, I would suggest let's clearly distinguish between cases that may cause errors and not.

Code after:
        data = [{"image": input_file}]

        test_transforms = Compose([LoadImaged(keys=["image"]), EnsureChannelFirstd(keys=["image"])])

        post_transforms = Compose(
            [
                SaveImaged(
                    keys=save_keys,
                    meta_keys="image_meta_dict",
                    output_dir=self.output_dir,
                    output_postfix="seg",
                    savepath_in_metadict=True,
                ),
                WriteFileMappingd(keys=write_keys, mapping_file_path=self.mapping_file_path),
            ]
        )

        dataset = Dataset(data=data, transform=test_transforms)
        dataloader = DataLoader(dataset, batch_size=1)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = UNet(spatial_dims=3, in_channels=1, out_channels=2, channels=(16, 32), strides=(2,)).to(device)
        model.eval()

        with torch.no_grad():
            for batch_data in dataloader:
                test_inputs = batch_data["image"].to(device)
                roi_size = (64, 64, 64)
                sw_batch_size = 2
                batch_data["seg"] = sliding_window_inference(test_inputs, roi_size, sw_batch_size, model)
                batch_data = [post_transforms(i) for i in decollate_batch(batch_data)]

        return input_file, output_file

    @parameterized.expand(SUCCESS_CASES)
    def test_successful_mapping_filed(self, save_keys, write_keys):
        input_file, output_file = self.run_test(save_keys, write_keys)
        self.assertTrue(os.path.exists(self.mapping_file_path))
        with open(self.mapping_file_path) as f:
            mapping_data = json.load(f)
        self.assertEqual(len(mapping_data), len(write_keys))
        for entry in mapping_data:
            self.assertEqual(entry["input"], input_file)
            self.assertEqual(entry["output"], output_file)

    @parameterized.expand(FAILURE_CASES)
    def test_failure_mapping_filed(self, save_keys, write_keys):
        with self.assertRaises(RuntimeError) as cm:
            self.run_test(save_keys, write_keys)

        cause_exception = cm.exception.__cause__
        self.assertIsInstance(cause_exception, KeyError)