"""
Tests for the Feature Provider
"""
import numpy as np
import os
import sys
import unittest
import warnings
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
sys.path.insert(0, path)
import senses.dataproviders.featureprovider as fp # pylint: disable=locally-disabled, import-error

class TestFeatureProvider(unittest.TestCase):
    def setUp(self):
        warnings.simplefilter("ignore", ResourceWarning)
        self.root = os.path.abspath("test_data_directory")
        self.sample_rate = 24_000
        self.nchannels = 1
        self.bytewidth = 2
        self.provider = fp.FeatureProvider(self.root, sample_rate=self.sample_rate, nchannels=self.nchannels, bytewidth=self.bytewidth)

    def _reset(self):
        """
        Resets the iterator.
        """
        self.provider = fp.FeatureProvider(self.root, sample_rate=self.sample_rate, nchannels=self.nchannels, bytewidth=self.bytewidth)

    def _label_fn(self, fpath):
        """
        Returns 0 if fpath contains 'babies', otherwise 1.
        """
        if "babies" in fpath:
            return 0
        else:
            return 1

    def test_generate_one_fft(self):
        """
        Tests yielding a single FFT of 45ms.
        """
        ffts = [(f, label) for f, label in self.provider.generate_n_ffts(n=1, ms=45, label_fn=self._label_fn)]
        self.assertEqual(len(ffts), 1)
        fft, _label = ffts[0]
        self.assertEqual(fft.shape, (541,))

    def test_generate_more_than_one_fft(self):
        """
        Tests yielding more than one FFT of 34ms each.