#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (c) 2017 Satpy developers
#
# This file is part of satpy.
#
# satpy is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# satpy is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with
# satpy. If not, see <http://www.gnu.org/licenses/>.
"""Tests for the geotiff writer."""
import unittest
from unittest import mock
import numpy as np
[docs]class TestGeoTIFFWriter(unittest.TestCase):
"""Test the GeoTIFF Writer class."""
[docs] def setUp(self):
"""Create temporary directory to save files to."""
import tempfile
self.base_dir = tempfile.mkdtemp()
[docs] def tearDown(self):
"""Remove the temporary directory created for a test."""
try:
import shutil
shutil.rmtree(self.base_dir, ignore_errors=True)
except OSError:
pass
def _get_test_datasets(self):
"""Create a single test dataset."""
from datetime import datetime
import dask.array as da
import xarray as xr
ds1 = xr.DataArray(
da.zeros((100, 200), chunks=50),
dims=('y', 'x'),
attrs={'name': 'test',
'start_time': datetime.utcnow()}
)
return [ds1]
[docs] def test_init(self):
"""Test creating the writer with no arguments."""
from satpy.writers.geotiff import GeoTIFFWriter
GeoTIFFWriter()
[docs] def test_simple_write(self):
"""Test basic writer operation."""
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir)
w.save_datasets(datasets)
[docs] def test_simple_delayed_write(self):
"""Test writing can be delayed."""
import dask.array as da
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir)
# when we switch to rio_save on XRImage then this will be sources
# and targets
res = w.save_datasets(datasets, compute=False)
# this will fail if rasterio isn't installed
self.assertIsInstance(res, tuple)
# two lists, sources and destinations
self.assertEqual(len(res), 2)
self.assertIsInstance(res[0], list)
self.assertIsInstance(res[1], list)
self.assertIsInstance(res[0][0], da.Array)
da.store(res[0], res[1])
for target in res[1]:
if hasattr(target, 'close'):
target.close()
[docs] def test_colormap_write(self):
"""Test writing an image with a colormap."""
from trollimage.colormap import spectral
from trollimage.xrimage import XRImage
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir)
# we'd have to customize enhancements to test this through
# save_datasets. We'll use `save_image` as a workaround.
img = XRImage(datasets[0])
img.palettize(spectral)
w.save_image(img, keep_palette=True)
[docs] def test_float_write(self):
"""Test that geotiffs can be written as floats.
NOTE: Does not actually check that the output is floats.
"""
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir,
enhance=False,
dtype=np.float32)
w.save_datasets(datasets)
[docs] def test_dtype_for_enhance_false(self):
"""Test that dtype of dataset is used if parameters enhance=False and dtype=None."""
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir, enhance=False)
with mock.patch('satpy.writers.XRImage.save') as save_method:
save_method.return_value = None
w.save_datasets(datasets, compute=False)
self.assertEqual(save_method.call_args[1]['dtype'], np.float64)
[docs] def test_dtype_for_enhance_false_and_given_dtype(self):
"""Test that dtype of dataset is used if enhance=False and dtype=uint8."""
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir, enhance=False, dtype=np.uint8)
with mock.patch('satpy.writers.XRImage.save') as save_method:
save_method.return_value = None
w.save_datasets(datasets, compute=False)
self.assertEqual(save_method.call_args[1]['dtype'], np.uint8)
[docs] def test_fill_value_from_config(self):
"""Test fill_value coming from the writer config."""
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir)
w.info['fill_value'] = 128
with mock.patch('satpy.writers.XRImage.save') as save_method:
save_method.return_value = None
w.save_datasets(datasets, compute=False)
self.assertEqual(save_method.call_args[1]['fill_value'], 128)
[docs] def test_scale_offset(self):
"""Test tags being added."""
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(tags={'test1': 1}, base_dir=self.base_dir)
w.info['fill_value'] = 128
with mock.patch('satpy.writers.XRImage.save') as save_method:
save_method.return_value = None
w.save_datasets(datasets, tags={'test2': 2}, compute=False, include_scale_offset=True)
called_include = save_method.call_args[1]['include_scale_offset_tags']
self.assertTrue(called_include)
[docs] def test_tiled_value_from_config(self):
"""Test tiled value coming from the writer config."""
from satpy.writers.geotiff import GeoTIFFWriter
datasets = self._get_test_datasets()
w = GeoTIFFWriter(base_dir=self.base_dir)
with mock.patch('satpy.writers.XRImage.save') as save_method:
save_method.return_value = None
w.save_datasets(datasets, compute=False)
self.assertEqual(save_method.call_args[1]['tiled'], True)