test_attribute.py

# coding=utf-8
from unittest import TestCase
import os
from tempfile import mktemp

import dummydata

from earthdiagnostics.diagnostic import DiagnosticVariableOption
from earthdiagnostics.box import Box
from earthdiagnostics.general.attribute import Attribute
from mock import Mock, patch

from earthdiagnostics.modelingrealm import ModelingRealms


class TestAttribute(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()
        self.box.min_depth = 0
        self.box.max_depth = 100
        self.var_file = mktemp('.nc')

    def tearDown(self):
        if os.path.exists(self.var_file):
            os.remove(self.var_file)

    def fake_parse(self, value):
        return value

    @patch.object(DiagnosticVariableOption, 'parse', fake_parse)
    def test_generate_jobs(self):

        jobs = Attribute.generate_jobs(self.diags, ['diagnostic', 'atmos', 'var', 'att', 'value'])
        self.assertEqual(len(jobs), 2)
        self.assertEqual(jobs[0], Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', '',
                                            'att', 'value'))
        self.assertEqual(jobs[1], Attribute(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos, 'var', '',
                                            'att', 'value'))

        jobs = Attribute.generate_jobs(self.diags, ['diagnostic', 'seaice', 'var', 'att', 'value', 'grid'])
        self.assertEqual(len(jobs), 2)
        self.assertEqual(jobs[0], Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce, 'var', 'grid',
                                            'att', 'value'))
        self.assertEqual(jobs[1], Attribute(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce, 'var', 'grid',
                                            'att', 'value'))

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

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

    def test_str(self):
        mixed = Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'var', 'grid', 'att', 'value')
        self.assertEqual(str(mixed),
                         'Write attributte output Startdate: 20010101 Member: 0 Chunk: 0 Variable: atmos:var '
                         'Attributte: att:value Grid: grid')

    def test_compute(self):
        dummydata.model3.Model3(oname=self.var_file, var='ta', start_year=2000, stop_year=2000, method='constant',
                                constant=1)

        diag = Attribute(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'ta', 'grid', 'att', 'value')

        diag.variable_file = Mock()
        diag.variable_file.local_file = self.var_file
        diag.corrected = Mock()
        diag.compute()
        diag.corrected.set_local_file.assert_called_once()