def test_test_types(datamodule: MVTecAD2) -> None:
        """Test that the datamodule can handle different test types."""
        # Test public test set
        public_loader = datamodule.test_dataloader(test_type=TestType.PUBLIC)
        assert public_loader is not None
        batch = next(iter(public_loader))
        assert batch.image.shape == (4, 3, 256, 256)

        # Test private test set
        private_loader = datamodule.test_dataloader(test_type=TestType.PRIVATE)
        assert private_loader is not None
        batch = next(iter(private_loader))
        assert batch.image.shape == (4, 3, 256, 256)

        # Test private mixed test set
        mixed_loader = datamodule.test_dataloader(test_type=TestType.PRIVATE_MIXED)
        assert mixed_loader is not None
        batch = next(iter(mixed_loader))
        assert batch.image.shape == (4, 3, 256, 256)

        # Test invalid test type
        with pytest.raises(ValueError, match="'invalid' is not a valid TestType"):
            datamodule.test_dataloader(test_type="invalid")
Test that the datamodule can handle different test types.
tests/unit/data/datamodule/image/test_mvtecad2.py