def test_get_item_returns_correct_keys_and_shapes(subset: str, datamodule: AnomalibDataModule) -> None:
        """Test that the datamodule __getitem__ returns image, mask, label and boxes."""
        # Get the dataloader.
        dataloader = getattr(datamodule, f"{subset}_dataloader")()

        # Get the first batch.
        batch = next(iter(dataloader))

        # Check that the batch has the correct shape.
        assert batch.image.shape == (4, 3, 256, 256)
        assert batch.gt_label.shape == (4,)

        if dataloader.dataset.task in {"detection", "segmentation"}:
            assert batch.gt_mask.shape == (4, 256, 256)
Test that the datamodule __getitem__ returns image, mask, label and boxes.
tests/unit/data/datamodule/base/image.py