test_module.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.module import Module
from mock import Mock, patch

from earthdiagnostics.modelingrealm import ModelingRealms


class TestModule(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.varu_file = mktemp('.nc')
        self.varv_file = mktemp('.nc')

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

    def fake_parse(self, value):
        return value

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

        jobs = Module.generate_jobs(self.diags, ['diagnostic', 'atmos', 'varu', 'varv', 'varmodule'])
        self.assertEqual(len(jobs), 2)
        self.assertEqual(jobs[0], Module(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos,
                                         'varu', 'varv', 'varmodule', ''))
        self.assertEqual(jobs[1], Module(self.data_manager, '20010101', 0, 1, ModelingRealms.atmos,
                                         'varu', 'varv', 'varmodule', ''))

        jobs = Module.generate_jobs(self.diags, ['diagnostic', 'seaIce', 'varu', 'varv', 'varmodule', 'grid'])
        self.assertEqual(len(jobs), 2)
        self.assertEqual(jobs[0], Module(self.data_manager, '20010101', 0, 0, ModelingRealms.seaIce,
                                         'varu', 'varv', 'varmodule', 'grid'))
        self.assertEqual(jobs[1], Module(self.data_manager, '20010101', 0, 1, ModelingRealms.seaIce,
                                         'varu', 'varv', 'varmodule', 'grid'))

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

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

    def test_str(self):
        mixed = Module(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'varu', 'varv', 'varmodule', 'grid')
        self.assertEqual(str(mixed),
                         'Calculate module Startdate: 20010101 Member: 0 Chunk: 0 '
                         'Variables: atmos:varu,varv,varmodule Grid: grid')

    def test_compute(self):
        dummydata.model2.Model2(oname=self.varu_file, var='ua', start_year=2000, stop_year=2000, method='constant',
                                constant=1)

        dummydata.model2.Model2(oname=self.varv_file, var='va', start_year=2000, stop_year=2000, method='constant',
                                constant=1)

        diag = Module(self.data_manager, '20010101', 0, 0, ModelingRealms.atmos, 'ua', 'va', 'varmodule', 'grid')

        diag.component_u_file = Mock()
        diag.component_u_file.local_file = self.varu_file

        diag.component_v_file = Mock()
        diag.component_v_file.local_file = self.varv_file
        diag.module_file = Mock()
        diag.compute()
        diag.module_file.set_local_file.assert_called_once()