test_preface.py

import os
import shutil
import tempfile

from unittest import TestCase
from mock import Mock
import mock

from earthdiagnostics.modelingrealm import ModelingRealms
from earthdiagnostics.frequency import Frequencies
from earthdiagnostics.data_convention import PrefaceConvention


class TestPrefaceConvention(TestCase):
    def setUp(self):
        self.tmp_dir = tempfile.mkdtemp()
        os.mkdir(os.path.join(self.tmp_dir, "expid"))
        self.config = Mock()
        self.config.data_dir = self.tmp_dir

        self.config.experiment.experiment_name = "experiment_name"
        self.config.experiment.expid = "expid"
        self.config.experiment.institute = "institute"
        self.config.experiment.model = "model"
        self.config.experiment.member_count_start = 0
        self.config.experiment.chunk_size = 1
        self.config.experiment.calendar = "standard"

        self.config.cmor.activity = "activity"
        self.config.cmor.append_startdate = False
        self.config.cmor.initialization_number = 1
        self.config.cmor.version = "version"

        self.convention = PrefaceConvention("name", self.config)

    def tearDown(self):
        shutil.rmtree(self.tmp_dir)

    def test_get_startdate_path(self):
        self.assertEqual(
            self.convention.get_startdate_path("19900101"),
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S19900101",
            ),
        )

    def test_experiment_name(self):
        self.assertEqual(
            self.convention.experiment_name("19900101"), "experiment_name"
        )

    def test_experiment_name_append(self):
        self.config.cmor.append_startdate = True
        self.assertEqual(
            self.convention.experiment_name("19900101"),
            "experiment_nameS19900101",
        )

    def test_get_cmor_folder_path(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_cmor_folder_path(
            "19900101", 1, ModelingRealms.ocean, "var", "mon", None, cmor_var
        )
        self.assertEqual(
            file_path,
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S19900101/mon/ocean/var/r2i1p1/version",
            ),
        )

    def test_get_cmor_folder_path_atmos(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_cmor_folder_path(
            "19900101", 1, ModelingRealms.atmos, "var", "mon", None, cmor_var
        )
        self.assertEqual(
            file_path,
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S19900101/mon/atmos/var/r2i1p1/version",
            ),
        )

    def test_get_cmor_folder_path_custom_grid(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_cmor_folder_path(
            "19900101", 1, ModelingRealms.ocean, "var", "mon", "grid", cmor_var
        )
        self.assertEqual(
            file_path,
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S19900101/mon/ocean/var/grid/r2i1p1/version",
            ),
        )

    def test_get_cmor_folder_path_no_cmor(self):
        frequency = Mock()
        frequency.__str__ = Mock()
        frequency.__str__.return_value = "mon"
        frequency.frequency = "mon"
        file_path = self.convention.get_cmor_folder_path(
            "19900101",
            1,
            ModelingRealms.ocean,
            "var",
            Frequencies.monthly,
            None,
            None,
        )
        self.assertEqual(
            file_path,
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S19900101/mon/ocean/var/r2i1p1/version",
            ),
        )

    def test_get_file_path_no_version(self):
        self.config.cmor.version = ""
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_cmor_folder_path(
            "19900101",
            1,
            ModelingRealms.ocean,
            "var",
            Frequencies.monthly,
            None,
            None,
        )
        self.assertEqual(
            file_path,
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S19900101/mon/ocean/var/r2i1p1",
            ),
        )

    def test_get_filename(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_file_name(
            "19900101",
            1,
            ModelingRealms.ocean,
            "var",
            cmor_var,
            Frequencies.monthly,
            1,
            None,
            None,
            None,
        )
        self.assertEqual(
            file_path,
            "var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc",
        )

    def test_get_filename_daily(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_file_name(
            "19900101",
            1,
            ModelingRealms.ocean,
            "var",
            cmor_var,
            Frequencies.daily,
            1,
            None,
            None,
            None,
        )
        self.assertEqual(
            file_path,
            "var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc",
        )

    def test_get_filename_atmos(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Amon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_file_name(
            "19900101",
            1,
            ModelingRealms.atmos,
            "var",
            cmor_var,
            Frequencies.monthly,
            1,
            None,
            None,
            None,
        )
        self.assertEqual(
            file_path,
            "var_Amon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc",
        )

    def test_get_filename_grid(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_file_name(
            "19900101",
            1,
            ModelingRealms.ocean,
            "var",
            cmor_var,
            Frequencies.monthly,
            1,
            None,
            None,
            "grid",
        )
        self.assertEqual(
            file_path,
            "var_Omon_model_experiment_name_S19900101_r2i1p1_199001_199001.nc",
        )

    def test_get_filename_year(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_file_name(
            "19900101",
            1,
            ModelingRealms.ocean,
            "var",
            cmor_var,
            Frequencies.yearly,
            None,
            1990,
            None,
            None,
        )
        self.assertEqual(
            file_path,
            "var_Omon_model_experiment_name_S19900101_r2i1p1_1990.nc",
        )

    def test_get_filename_date_str(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        file_path = self.convention.get_file_name(
            "19900101",
            1,
            ModelingRealms.ocean,
            "var",
            cmor_var,
            Frequencies.monthly,
            None,
            None,
            "date_str",
            None,
        )
        self.assertEqual(
            file_path,
            "var_Omon_model_experiment_name_S19900101_r2i1p1_date_str.nc",
        )

    def test_get_filename_no_date_info(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon

        with self.assertRaises(ValueError):
            self.convention.get_file_name(
                "19900101",
                1,
                ModelingRealms.ocean,
                "var",
                cmor_var,
                Frequencies.monthly,
                None,
                None,
                None,
                None,
            )

    @mock.patch("os.path.isfile")
    def test_is_cmorized(self, mock_is_file):
        mock_is_file.return_value = True
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        self.config.var_manager.get_variable.return_value = cmor_var
        self.config.cmor.min_cmorized_vars = 1
        os.makedirs(
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S20000101/mon/ocean/var/r2i1p1/version",
            )
        )
        self.assertTrue(
            self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean)
        )

    @mock.patch("os.path.isfile")
    def test_is_not_cmorized(self, mock_is_file):
        mock_is_file.return_value = False
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        self.config.var_manager.get_variable.return_value = cmor_var
        self.config.cmor.min_cmorized_vars = 1
        os.makedirs(
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S20000101/mon/ocean/var/r2i1p1/version",
            )
        )
        self.assertFalse(
            self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean)
        )

    def test_is_cmorized_false_not_member_folder(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        self.config.var_manager.get_variable.return_value = cmor_var
        self.config.cmor.min_cmorized_vars = 1
        os.makedirs(
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S20000101/mon/ocean/var/",
            )
        )
        self.assertFalse(
            self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean)
        )

    @mock.patch("os.path.isfile")
    def test_is_cmorized_not_enough_vars(self, mock_is_file):
        mock_is_file.return_value = True
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        self.config.var_manager.get_variable.return_value = cmor_var
        self.config.cmor.min_cmorized_vars = 2
        os.makedirs(
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S20000101/mon/ocean/var/r2i1p1/version",
            )
        )
        self.assertFalse(
            self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean)
        )

    def test_is_cmorized_not_domain_folder(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        self.config.var_manager.get_variable.return_value = cmor_var
        self.config.cmor.min_cmorized_vars = 2
        os.makedirs(
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S20000101/mon/r2i1p1/version",
            )
        )
        self.assertFalse(
            self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean)
        )

    def test_is_cmorized_not_freq_folder(self):
        cmor_var = Mock()
        omon = Mock()
        omon.name = "Omon"
        cmor_var.get_table.return_value = omon
        self.config.var_manager.get_variable.return_value = cmor_var
        self.config.cmor.min_cmorized_vars = 2
        os.makedirs(
            os.path.join(
                self.tmp_dir,
                "expid/original_files/cmorfiles/institute/experiment_name/"
                "S20000101",
            )
        )
        self.assertFalse(
            self.convention.is_cmorized("20000101", 1, 1, ModelingRealms.ocean)
        )

    @mock.patch(
        "earthdiagnostics.data_convention.PrefaceConvention.create_link"
    )
    def test_create_links(self, mock_create_link):
        member_path = os.path.join(
            self.tmp_dir,
            "expid/original_files/cmorfiles/institute/experiment_name/"
            "S20010101/mon/ocean/var/r2i1p1",
        )
        os.makedirs(member_path)
        tempfile.mkstemp(dir=member_path)
        self.convention.create_links("20010101", 1)
        mock_create_link.assert_called()

    @mock.patch(
        "earthdiagnostics.data_convention.PrefaceConvention.create_link"
    )
    def test_create_links_member_not_found(self, mock_create_link):
        member_path = os.path.join(
            self.tmp_dir,
            "expid/original_files/cmorfiles/institute/experiment_name/"
            "S20010101/mon/ocean/var/r1i1p1",
        )
        os.makedirs(member_path)
        tempfile.mkstemp(dir=member_path)
        self.convention.create_links("20010101", 1)
        mock_create_link.assert_not_called()

    @mock.patch(
        "earthdiagnostics.data_convention.PrefaceConvention.create_link"
    )
    def test_create_links_with_grid(self, mock_create_link):
        member_path = os.path.join(
            self.tmp_dir,
            "expid/original_files/cmorfiles/institute/experiment_name/"
            "S20010101/mon/ocean/var/r2i1p1/grid",
        )
        os.makedirs(member_path)
        tempfile.mkstemp(dir=member_path)
        self.convention.create_links("20010101", 1)
        mock_create_link.assert_called()