test_psi.py

# coding=utf-8
from unittest import TestCase
from mock import Mock, patch

from earthdiagnostics.constants import Basins
from earthdiagnostics.ocean.psi import Psi
from earthdiagnostics.utils import Utils


class TestPsi(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),
        )

    @staticmethod
    def fake_get(basin):
        return None

    @patch.object(Utils, "get_mask", fake_get)
    def test_generate_jobs(self):
        jobs = Psi.generate_jobs(self.diags, ["diagnostic"])
        self.assertEqual(len(jobs), 2)
        self.assertEqual(
            jobs[0],
            Psi(self.data_manager, "20010101", 0, 0, {Basins().Global: None}),
        )
        self.assertEqual(
            jobs[1],
            Psi(self.data_manager, "20010101", 0, 1, {Basins().Global: None}),
        )

        with self.assertRaises(Exception):
            Psi.generate_jobs(self.diags, ["diagnostic", "extra"])

    def test_str(self):
        psi = Psi(self.data_manager, "20000101", 1, 1, {Basins().Global: None})
        self.assertEqual(
            str(psi),
            "PSI Startdate: 20000101 Member: 1 Chunk: 1 Basins:Global",
        )