From 6e8401909ff9983bdda48c451eb82e21deb7cfa3 Mon Sep 17 00:00:00 2001
From: vincentsarago <vincent.sarago@gmail.com>
Date: Fri, 4 Oct 2024 12:42:34 +0200
Subject: [PATCH] refactor benchmark

---
 tests/benchmarks/benchmarks.py | 32 ++++++++++++++++----------------
 tests/conftest.py              | 33 +++++++++++++++++++++++++++------
 2 files changed, 43 insertions(+), 22 deletions(-)

diff --git a/tests/benchmarks/benchmarks.py b/tests/benchmarks/benchmarks.py
index 8c1559af..f0221aa3 100644
--- a/tests/benchmarks/benchmarks.py
+++ b/tests/benchmarks/benchmarks.py
@@ -71,9 +71,8 @@
 def read_tile(dst, tile):
     """Benchmark rio-tiler.utils._tile_read."""
     # We make sure to not store things in cache.
-    with rasterio.Env(GDAL_CACHEMAX=0, NUM_THREADS="all"):
-        with Reader(None, dataset=dst) as src:
-            return src.tile(*tile)
+    with Reader(None, dataset=dst) as src:
+        return src.tile(*tile)
 
 
 data_types = list(dtype_ranges.keys())
@@ -93,16 +92,17 @@ def test_tile(
     tile = benchmark_tiles[dataset_name][tile_name]
 
     dst_info = datasets[dataset_name]
-    with MemoryFile(
-        dataset_fixture(
-            crs=dst_info["crs"],
-            bounds=list(dst_info["bounds"]),
-            dtype=data_type,
-            nodata_type=nodata_type,
-            width=256,
-            height=256,
-        )
-    ) as memfile:
-        with memfile.open() as dst:
-            img = benchmark(read_tile, dst, tile)
-            assert img.data.dtype == data_type
+    with rasterio.Env(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", NUM_THREADS="all"):
+        with MemoryFile(
+            dataset_fixture(
+                crs=dst_info["crs"],
+                bounds=list(dst_info["bounds"]),
+                dtype=data_type,
+                nodata_type=nodata_type,
+                width=4000,
+                height=4000,
+            )
+        ) as memfile:
+            with memfile.open() as dst:
+                img = benchmark(read_tile, dst, tile)
+                assert img.data.dtype == data_type
diff --git a/tests/conftest.py b/tests/conftest.py
index ada2ede2..e54d48e7 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -8,7 +8,10 @@
 import rasterio
 from rasterio.crs import CRS
 from rasterio.enums import ColorInterp
+from rasterio.enums import Resampling as ResamplingEnums
 from rasterio.io import MemoryFile
+from rasterio.rio.overview import get_maximum_overview_level
+from rasterio.shutil import copy
 from rasterio.transform import from_bounds
 
 with rasterio.Env() as env:
@@ -33,16 +36,18 @@ def _dataset(
         width: int = 256,
         height: int = 256,
     ):
+        max_value = 127 if dtype == "int8" else 255
+
         # Data
         arr = numpy.zeros((nband, height, width), dtype=dtype) + 1
+        arr[:, range(height), range(width)] = max_value
+        arr[:, range(height - 1, 0, -1), range(width - 1)] = max_value
+        arr[:, :, width // 2] = max_value
+        arr[:, height // 2, :] = max_value
         arr[:, 0:128, 0:128] = 0
 
         # Mask/Alpha
-        if dtype == "int8":
-            mask = numpy.zeros((1, height, width), dtype=dtype) + 127
-        else:
-            mask = numpy.zeros((1, height, width), dtype=dtype) + 255
-
+        mask = numpy.zeros((1, height, width), dtype=dtype) + max_value
         mask[:, 0:128, 0:128] = 0
 
         # Input Profile
@@ -89,6 +94,22 @@ def _dataset(
                     if nodata_type == "mask":
                         mem.write_mask(mask[0])
 
-                return BytesIO(memfile.read())
+                    overview_level = get_maximum_overview_level(
+                        mem.width, mem.height, minsize=512
+                    )
+                    overviews = [2**j for j in range(1, overview_level + 1)]
+                    mem.build_overviews(overviews, ResamplingEnums.bilinear)
+
+                    cog_profile = {
+                        "interleave": "pixel",
+                        "compress": "DEFLATE",
+                        "tiled": True,
+                        "blockxsize": 512,
+                        "blockysize": 512,
+                    }
+
+                    with MemoryFile() as cogfile:
+                        copy(mem, cogfile.name, copy_src_overviews=True, **cog_profile)
+                        return BytesIO(cogfile.read())
 
     return _dataset