test_primavera.py

"""Tests for PRIMAVERA Convention"""
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 PrimaveraConvention


class TestPrimaveraConvention(TestCase):
    """Tests for PRIMAVERA convetion class"""

    def setUp(self):
        """Prepare tests"""
        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.config.cmor.default_ocean_grid = 'ocean_grid'
        self.config.cmor.default_atmos_grid = 'atmos_grid'

        self.convention = PrimaveraConvention('name', self.config)

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

    def test_get_startdate_path(self):
        """Test get startdate path"""
        self.assertEqual(self.convention.get_startdate_path('19900101'),
                         os.path.join(self.tmp_dir, 'expid/cmorfiles/activity/institute/model/experiment_name'))

    def test_experiment_name(self):
        """Test get expriment name"""
        self.assertEqual(self.convention.experiment_name('19900101'),
                         'experiment_name')

    def test_experiment_name_append(self):
        """Test get expriment name when appending startdate"""
        self.config.cmor.append_startdate = True
        self.assertEqual(self.convention.experiment_name('19900101'),
                         'experiment_name')

    def test_get_cmor_folder_path(self):
        """Test get cmor foilder path"""
        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/cmorfiles/activity/institute/model/experiment_name/'
                                                    'r2i1p1f1/Omon/var/ocean_grid/version'))

    def test_get_cmor_folder_path_no_cmor_var(self):
        """Test get cmor folder path when not passing cmor_var"""
        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/cmorfiles/activity/institute/model/experiment_name/'
                                                    'r2i1p1f1/Omon/var/ocean_grid/version'))

    def test_get_cmor_folder_path_atmos(self):
        """Test get cmor foilder path for the atmos"""
        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/cmorfiles/activity/institute/model/experiment_name/'
                                                    'r2i1p1f1/Omon/var/atmos_grid/version'))

    def test_get_cmor_folder_path_custom_grid(self):
        """Test get cmor foilder path for a custom grid"""
        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/cmorfiles/activity/institute/model/experiment_name/'
                                                    'r2i1p1f1/Omon/var/grid/version'))

    def test_get_cmor_folder_path_no_cmor(self):
        """Test get cmor folder path with no cmor_var"""
        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/cmorfiles/activity/institute/model/experiment_name/'
                                                    'r2i1p1f1/Omon/var/ocean_grid/version'))

    def test_get_file_path_no_version_primavera(self):
        """Test get cmor folder path with no version"""
        self.config.cmor.version = ''
        cmor_var = Mock()
        omon = Mock()
        omon.name = 'Omon'
        cmor_var.get_table.return_value = omon
        with self.assertRaises(ValueError):
            self.convention.get_cmor_folder_path('19900101', 1, ModelingRealms.ocean, 'var', 'mon', 'grid', cmor_var)

    def test_get_filename(self):
        """Test get_filename"""
        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_r2i1p1f1_ocean_grid_199001-199001.nc')

    def test_get_filename_with_startdate(self):
        """Test get_filename with startdate"""
        self.config.cmor.append_startdate = True
        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-r2i1p1f1_ocean_grid_199001-199001.nc')

    def test_get_filename_no_cmor_var(self):
        """Test get_filename not passing cmor_var"""
        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', None,
                                                  Frequencies.monthly, 1, None, None, None)
        self.assertEqual(file_path,
                         'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_199001-199001.nc')

    def test_get_filename_daily(self):
        """Test get_filename for daily frequency"""
        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_r2i1p1f1_ocean_grid_19900101-19900131.nc')

    def test_get_filename_6hourly(self):
        """Test get_filename for 6hourly files"""
        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.six_hourly, 1,
                                                  None, None, None)
        self.assertEqual(file_path,
                         'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_199001010000-199001311800.nc')

    def test_get_filename_3hourly(self):
        """Test get_filename for 3hourly files"""
        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.three_hourly, 1,
                                                  None, None, None)
        self.assertEqual(file_path,
                         'var_Omon_model_experiment_name_r2i1p1f1_ocean_grid_199001010000-199001312100.nc')

    def test_get_filename_atmos(self):
        """Test get_filename for atmos"""
        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.atmos, 'var', cmor_var,
                                                  Frequencies.monthly, 1, None, None, None)
        self.assertEqual(file_path,
                         'var_Omon_model_experiment_name_r2i1p1f1_atmos_grid_199001-199001.nc')

    def test_get_filename_grid(self):
        """Test get_filename for a custom grid"""
        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_r2i1p1f1_grid_199001-199001.nc')

    def test_get_filename_year(self):
        """Test get_filename for a whole year"""
        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_r2i1p1f1_ocean_grid_1990.nc')

    def test_get_filename_date_Str(self):
        """Test get_filename passing date_Str"""
        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_r2i1p1f1_ocean_grid_date_str.nc')

    def test_get_filename_no_date_info(self):
        """Test get_filename with no date info raises ValueError"""
        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):
        """Test is cmorized"""
        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/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/var'))
        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):
        """Test is cmorized false"""
        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/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/var'))
        self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean))

    def test_is_cmorized_false_not_member_folder(self):
        """Test is cmorized false bacause ther is no member folder"""
        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/cmorfiles/activity/institute/model/experiment_name/'))
        self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean))

    def test_is_cmorized_false_not_table_folder(self):
        """Test is cmorized false bacause ther is no table folder"""
        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/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1'))
        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):
        """Test is cmorized false because thera are not eniouch variables"""
        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/cmorfiles/institute/model/experiment_name/S20000101/'
                                               'mon/ocean/var'))
        self.assertFalse(self.convention.is_cmorized('20000101', 1, 1, ModelingRealms.ocean))

    @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link')
    def test_create_links_primavera(self, mock_create_link):
        """Test create links"""
        member_path = os.path.join(self.tmp_dir,
                                   'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/var/gn')
        os.makedirs(member_path)
        self.config.var_manager.tables = {'Omon': Mock()}
        tempfile.mkstemp(dir=member_path)
        self.convention.create_links('20010101', 1)
        mock_create_link.assert_called()

    @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link')
    def test_create_links_with_version_primavera(self, mock_create_link):
        """Test create links with version"""
        member_path = os.path.join(self.tmp_dir,
                                   'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/'
                                   'var/gn/version')
        os.makedirs(member_path)
        self.config.var_manager.tables = {'Omon': Mock()}
        tempfile.mkstemp(dir=member_path)
        self.convention.create_links('20010101', 1)
        mock_create_link.assert_called()

    @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link')
    def test_create_links_with_version_primavera_no_member(self, mock_create_link):
        """Test create links with version full startdate"""
        member_path = os.path.join(self.tmp_dir,
                                   'expid/cmorfiles/activity/institute/model/experiment_name/r2i1p1f1/Omon/'
                                   'var/gn/version')
        os.makedirs(member_path)
        self.config.var_manager.tables = {'Omon': Mock()}
        tempfile.mkstemp(dir=member_path)
        self.convention.create_links('20010101')
        mock_create_link.assert_called()

    @mock.patch('earthdiagnostics.data_convention.PrimaveraConvention.create_link')
    def test_create_links_member_not_found_primavera(self, mock_create_link):
        """Test create links when the member can not be found"""
        member_path = os.path.join(self.tmp_dir,
                                   'expid/cmorfiles/activity/institute/model/experiment_name/r1i1p1f1/Omon/var/gn')
        os.makedirs(member_path)
        self.config.var_manager.tables = {'Omon': Mock()}
        tempfile.mkstemp(dir=member_path)
        self.convention.create_links('20010101', 1)
        mock_create_link.assert_not_called()