Source code for satpy.tests.writer_tests.test_geotiff

#!/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_tags(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) called_tags = save_method.call_args[1]['tags'] self.assertDictEqual(called_tags, {'test1': 1, 'test2': 2})
[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)