diff --git a/test/unit_tests/conftest.py b/test/unit_tests/conftest.py
index de5aa0a..b904c76 100644
--- a/test/unit_tests/conftest.py
+++ b/test/unit_tests/conftest.py
@@ -100,9 +100,9 @@ def array_to_image(array):
     return Image.fromarray(np.uint8(array * 255), mode="RGB")
 
 
-@pytest.fixture(params=[{'depth': 3}, {'width': 10}, {'height': 15}])
+@pytest.fixture(params=[{"width": 10, "height": 15, "depth": 3}, {"width": 150, "height": 100, "depth": 3}])
 def input_size(request):
-    return request.param.get("width", 10), request.param.get("height", 15), request.param.get("depth", 3)
+    return itemgetter("width", "height", "depth")(request.param)
 
 
 @pytest.fixture
diff --git a/test/unit_tests/image_extractor_test.py b/test/unit_tests/image_extractor_test.py
index ffc311a..b07d711 100644
--- a/test/unit_tests/image_extractor_test.py
+++ b/test/unit_tests/image_extractor_test.py
@@ -13,8 +13,8 @@ def test_image_extractor_mock(image_extractor, images):
 
 
 @pytest.mark.parametrize("extractor_type", ["parsable_pdf"])
-@pytest.mark.parametrize("batch_size", [128])
-@pytest.mark.parametrize("input_size", [{"depth": 3, "width": 100, "height": 90}], indirect=["input_size"])
+@pytest.mark.parametrize("batch_size", [0, 1, 2, 64])
+@pytest.mark.parametrize("input_size", [{"depth": 3, "width": 170, "height": 220}], indirect=["input_size"])
 def test_parsable_pdf_image_extractor(image_extractor, pdf, images, metadata, input_size):
     images_extracted, metadata_extracted = map(list, extract_images_from_pdf(pdf, image_extractor))
     assert np.allclose(images_to_batch_tensor(images_extracted), images_to_batch_tensor(images))