test_heatcontent.py 1.81 KB
Newer Older
# coding=utf-8
from unittest import TestCase

from earthdiagnostics.box import Box
from earthdiagnostics.constants import Basins
from earthdiagnostics.ocean.heatcontent import HeatContent
from mock import Mock, patch


def _get_levels_from_meters_mock(cls, box):
    return 20, 10


class TestHeatContent(TestCase):

    def setUp(self):
        self.data_manager = Mock()

        self.diags = Mock()
        self.diags.model_version = 'model_version'
        self.diags.config.experiment.get_chunk_list.return_value = (('20010101', 0, 0), ('20010101', 0, 1))

        self.box = Box(True)
        self.box.min_depth = 0
        self.box.max_depth = 100

    @patch('earthdiagnostics.ocean.heatcontent.HeatContent._get_levels_from_meters')
    def test_generate_jobs(self, levels_mock):
        levels_mock.return_value = (1, 20)
        jobs = HeatContent.generate_jobs(self.diags, ['diagnostic', 'Global', '-1', '0', '100'])
        self.assertEqual(len(jobs), 2)
        self.assertEqual(jobs[0], HeatContent(self.data_manager, '20010101', 0, 0, Basins().Global, -1,
                                              self.box, 0, 0))
        self.assertEqual(jobs[1], HeatContent(self.data_manager, '20010101', 0, 1, Basins().Global, -1,
                                              self.box, 0, 0))

        with self.assertRaises(Exception):
            HeatContent.generate_jobs(self.diags, ['diagnostic'])

        with self.assertRaises(Exception):
            HeatContent.generate_jobs(self.diags, ['diagnostic', '0', '0', '0', '0', '0', '0', '0'])

    def test_str(self):
        diag = HeatContent(self.data_manager, '20010101', 0, 0, Basins().Global, -1, self.box, 1, 20)
        self.assertEquals(str(diag), 'Heat content Startdate: 20010101 Member: 0 Chunk: 0 Mixed layer: -1 Box: 0-100m '
                                     'Basin: Global')