# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import random
import string
import tempfile

import numpy as np
from PIL import Image

from mmselfsup.datasets.utils import check_integrity, rm_suffix, to_numpy


def test_to_numpy():
    pil_img = Image.open(osp.join(osp.dirname(__file__), '../data/color.jpg'))
    np_img = to_numpy(pil_img)
    assert type(np_img) == np.ndarray
    if np_img.ndim < 3:
        assert np_img.shape[0] == 1
    elif np_img.ndim == 3:
        assert np_img.shape[0] == 3


def test_dataset_utils():
    # test rm_suffix
    assert rm_suffix('a.jpg') == 'a'
    assert rm_suffix('a.bak.jpg') == 'a.bak'
    assert rm_suffix('a.bak.jpg', suffix='.jpg') == 'a.bak'
    assert rm_suffix('a.bak.jpg', suffix='.bak.jpg') == 'a'

    # test check_integrity
    rand_file = ''.join(random.sample(string.ascii_letters, 10))
    assert not check_integrity(rand_file, md5=None)
    assert not check_integrity(rand_file, md5=2333)
    tmp_file = tempfile.NamedTemporaryFile()
    assert check_integrity(tmp_file.name, md5=None)
    assert not check_integrity(tmp_file.name, md5=2333)