test_geotiff.py (1939B)
1 import math 2 import unittest 3 from pathlib import Path 4 from tempfile import TemporaryDirectory 5 6 import numpy as np 7 8 from tem_loader.mean_resistivity import MapGrid 9 10 try: 11 from osgeo import gdal, osr 12 except ImportError: 13 gdal = None 14 15 if gdal is not None: 16 osr.UseExceptions() 17 from tem_loader.geotiff import NODATA, write_geotiff 18 19 20 @unittest.skipIf(gdal is None, 'GDAL Python bindings not available') 21 class WriteGeotiffTests(unittest.TestCase): 22 def test_round_trip(self): 23 grid = MapGrid(500000.0, 6200000.0, 50.0, 3, 2) 24 array = np.array([[10.0, 20.0, math.nan], [30.0, 40.0, 50.0]]) 25 srs = osr.SpatialReference() 26 srs.ImportFromEPSG(25832) 27 28 with TemporaryDirectory() as tmp: 29 path = Path(tmp) / 'map.tif' 30 write_geotiff(path, grid, array, srs.ExportToWkt()) 31 dataset = gdal.Open(str(path)) 32 band = dataset.GetRasterBand(1) 33 34 self.assertEqual((dataset.RasterXSize, dataset.RasterYSize), 35 (3, 2)) 36 self.assertEqual(dataset.GetGeoTransform(), 37 (500000.0, 50.0, 0.0, 6200000.0, 0.0, -50.0)) 38 self.assertEqual(band.DataType, gdal.GDT_Float32) 39 self.assertEqual(band.GetNoDataValue(), NODATA) 40 self.assertEqual(band.ReadAsArray().tolist(), 41 [[10.0, 20.0, NODATA], [30.0, 40.0, 50.0]]) 42 written = osr.SpatialReference(wkt=dataset.GetProjection()) 43 self.assertEqual(written.GetAuthorityCode(None), '25832') 44 dataset = None 45 46 def test_unwritable_path_raises_value_error(self): 47 grid = MapGrid(0.0, 10.0, 10.0, 1, 1) 48 with TemporaryDirectory() as tmp: 49 path = Path(tmp) / 'missing' / 'map.tif' 50 with self.assertRaises(ValueError): 51 write_geotiff(path, grid, np.ones((1, 1)), '') 52 53 54 if __name__ == '__main__': 55 unittest.main()