test_region_mean.py

# coding=utf-8
from unittest import TestCase

from mock import Mock, patch

from earthdiagnostics.box import Box
from earthdiagnostics.frequency import Frequencies
from earthdiagnostics.constants import Basins
from earthdiagnostics.diagnostic import (
    DiagnosticOptionError,
    DiagnosticVariableListOption,
)
from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.ocean.regionmean import RegionMean
from earthdiagnostics.utils import TempFile


class TestRegionMean(TestCase):
    def setUp(self):
        self.data_manager = Mock()
        self.diags = Mock()
        self.diags.config.experiment.get_chunk_list.return_value = (
            ("20010101", 0, 0),
            ("20010101", 0, 1),
        )
        self.diags.config.frequency = Frequencies.monthly

    def fake_parse(self, value):
        if not value:
            raise DiagnosticOptionError
        return value

    @staticmethod
    def fake_get():
        return "tempfile"

    @patch.object(DiagnosticVariableListOption, "parse", fake_parse)
    @patch.object(TempFile, "get", fake_get)
    def test_generate_jobs(self):
        box = Box()
        box.min_depth = -1
        box.max_depth = -1

        jobs = RegionMean.generate_jobs(
            self.diags, ["diagnostic", "ocean", ["var"]]
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "t",
                Frequencies.monthly,
            ),
        )
        self.assertEqual(
            jobs[1],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "t",
                Frequencies.monthly,
            ),
        )

        jobs = RegionMean.generate_jobs(
            self.diags, ["diagnostic", "ocean", ["var"], "", "U"]
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )
        self.assertEqual(
            jobs[1],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )

        jobs = RegionMean.generate_jobs(
            self.diags, ["diagnostic", "ocean", ["var"], "global", "U"]
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )
        self.assertEqual(
            jobs[1],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )

        box = Box()
        box.min_depth = 1.0
        box.max_depth = 10.0

        jobs = RegionMean.generate_jobs(
            self.diags,
            ["diagnostic", "ocean", ["var"], "global", "U", "1", "10"],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )
        self.assertEqual(
            jobs[1],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                box,
                True,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )

        jobs = RegionMean.generate_jobs(
            self.diags,
            [
                "diagnostic",
                "ocean",
                ["var"],
                "global",
                "U",
                "1",
                "10",
                "false",
            ],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                box,
                False,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )
        self.assertEqual(
            jobs[1],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                box,
                False,
                [Basins().Global],
                "u",
                Frequencies.monthly,
            ),
        )

        jobs = RegionMean.generate_jobs(
            self.diags,
            [
                "diagnostic",
                "ocean",
                ["var"],
                "global",
                "U",
                "1",
                "10",
                "false",
                "grid",
            ],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                box,
                False,
                [Basins().Global],
                "grid",
                Frequencies.monthly,
            ),
        )
        self.assertEqual(
            jobs[1],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                box,
                False,
                [Basins().Global],
                "grid",
                Frequencies.monthly,
            ),
        )

        jobs = RegionMean.generate_jobs(
            self.diags,
            [
                "diagnostic",
                "ocean",
                ["var"],
                "global",
                "U",
                "1",
                "10",
                "false",
                "grid",
                "day",
            ],
        )
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                0,
                ModelingRealms.ocean,
                "var",
                box,
                False,
                [Basins().Global],
                "grid",
                Frequencies.daily,
            ),
        )
        self.assertEqual(
            jobs[1],
            RegionMean(
                self.data_manager,
                "20010101",
                0,
                1,
                ModelingRealms.ocean,
                "var",
                box,
                False,
                [Basins().Global],
                "grid",
                Frequencies.daily,
            ),
        )

        with self.assertRaises(DiagnosticOptionError):
            RegionMean.generate_jobs(self.diags, ["diagnostic"])

        with self.assertRaises(DiagnosticOptionError):
            RegionMean.generate_jobs(
                self.diags,
                [
                    "diagnostic",
                    "ocean",
                    ["var"],
                    "global",
                    "U",
                    "1",
                    "10",
                    "false",
                    "grid",
                    "day",
                    "extra",
                ],
            )

    def test_str(self):
        box = Box()
        box.min_depth = 1
        box.max_depth = 10
        diag = RegionMean(
            self.data_manager,
            "20010101",
            0,
            0,
            ModelingRealms.ocean,
            "var",
            box,
            False,
            [Basins().Global],
            "T",
            Frequencies.monthly,
        )
        self.assertEqual(
            str(diag),
            "Region mean Startdate: 20010101 Member: 0 Chunk: 0 Variable: var "
            "Box: 1-10 Save 3D: False Grid point: T",
        )