
from __future__ import absolute_import, division, print_function, unicode_literals

import os.path

import pytest

from cfgrib import eccodes
from cfgrib import xarray_store

SAMPLE_DATA_FOLDER = os.path.join(os.path.dirname(__file__), 'sample-data')
TEST_DATA = os.path.join(SAMPLE_DATA_FOLDER, 'era5-levels-members.grib')
TEST_CORRUPTED = os.path.join(SAMPLE_DATA_FOLDER, 'era5-levels-corrupted.grib')


def test_open_dataset():
    res = xarray_store.open_dataset(TEST_DATA)

    assert res.attrs['GRIB_edition'] == 1

    var = res['t']
    assert var.attrs['GRIB_gridType'] == 'regular_ll'
    assert var.attrs['units'] == 'K'
    assert var.dims == \
        ('number', 'time', 'isobaricInhPa', 'latitude', 'longitude')

    assert var.mean() > 0.

    with pytest.warns(FutureWarning):
        xarray_store.open_dataset(TEST_DATA, filter_by_keys={'typeOfLevel': 'isobaricInhPa'})


def test_open_dataset_corrupted():
    res = xarray_store.open_dataset(TEST_CORRUPTED)

    assert res.attrs['GRIB_edition'] == 1
    assert len(res.data_vars) == 1

    with pytest.raises(eccodes.EcCodesError):
        xarray_store.open_dataset(TEST_CORRUPTED, backend_kwargs={'grib_errors': 'strict'})


def test_open_dataset_encode_cf_time():
    backend_kwargs = {'encode_cf': ('time',)}
    res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)

    assert res.attrs['GRIB_edition'] == 1
    assert res['t'].attrs['GRIB_gridType'] == 'regular_ll'
    assert res['t'].attrs['GRIB_units'] == 'K'
    assert res['t'].dims == ('number', 'time', 'level', 'values')

    assert res['t'].mean() > 0.


def test_open_dataset_encode_cf_vertical():
    backend_kwargs = {'encode_cf': ('vertical',)}
    res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)

    var = res['t']
    assert var.dims == ('number', 'dataDate', 'dataTime', 'isobaricInhPa', 'values')

    assert var.mean() > 0.


def test_open_dataset_encode_cf_geography():
    backend_kwargs = {'encode_cf': ('geography',)}
    res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)

    assert res.attrs['GRIB_edition'] == 1

    var = res['t']
    assert var.attrs['GRIB_gridType'] == 'regular_ll'
    assert var.attrs['GRIB_units'] == 'K'
    assert var.dims == ('number', 'dataDate', 'dataTime', 'level', 'latitude', 'longitude')

    assert var.mean() > 0.


def test_open_dataset_eccodes():
    res = xarray_store.open_dataset(TEST_DATA)

    assert res.attrs['GRIB_edition'] == 1

    var = res['t']
    assert var.attrs['GRIB_gridType'] == 'regular_ll'
    assert var.attrs['units'] == 'K'
    assert var.dims == ('number', 'time', 'isobaricInhPa', 'latitude', 'longitude')

    assert var.mean() > 0.

@@ -2691,28 +2691,175 @@
 %0A    assert var.mean() %3E 0.%0A
+%0A%0Adef test_open_datasets():%0A    res = xarray_store.open_datasets(TEST_DATA)%0A%0A    assert len(res) == 1%0A    assert res%5B0%5D.attrs%5B'GRIB_edition'%5D == 1%0A

Add test for open_datasets.