Add tests for bands objects, and associated fixes

Author: TomDonoghue

fooof/bands.py

@@ -2,12 +2,10 @@
 
 from collections import OrderedDict
 
-import numpy as np
-
 ###################################################################################################
 ###################################################################################################
 
-class Bands(object):
+class Bands():
     """Class to hold definitions of oscillation bands.
 
     Attributes
@@ -16,19 +14,19 @@ class Bands(object):
         Dictionary of band definitions.
     """
 
-    def __init__(self, input_bands=None):
+    def __init__(self, input_bands={}):
         """Initialize the Bands object.
 
         Parameters
         ----------
-        input_bands : dict, optional (default = None)
+        input_bands : dict, optional
             A dictionary of oscillation bands to use.
         """
 
         self.bands = OrderedDict()
 
         for label, band_def in input_bands.items():
-            self.add_band(label)
+            self.add_band(label, band_def)
 
 
     def __getitem__(self, name):
@@ -46,7 +44,7 @@ def __getattr__(self, name):
 
     def __repr__(self):
 
-        return '\n'.join(['{:8} :  {:2} - {:2}  Hz'.format(*val) \
+        return '\n'.join(['{:8} :  {:2} - {:2}  Hz'.format(key, *val) \
             for key, val in self.bands.items()])
 
 
@@ -57,12 +55,14 @@ def __len__(self):
 
     @property
     def labels(self):
+        """Get the labels for all bands defined in the object."""
 
         return list(self.bands.keys())
 
 
     @property
     def n_bands(self):
+        """Get the number of bands defined in the object."""
 
         return len(self.bands)
 
@@ -78,8 +78,7 @@ def add_band(self, label, band_definition):
             The lower and upper frequency limit of the band, in Hz.
         """
 
-        # Check that band definition is properly formatted & add if so
-        _check_band(label, band_definition)
+        self._check_band(label, band_definition)
         self.bands[label] = band_definition
 
 
@@ -92,24 +91,24 @@ def remove_band(self, label):
             Band label to remove from band definitions.
         """
 
-        self.bands.pop(rm_band)
+        self.bands.pop(label)
 
 
     @staticmethod
     def _check_band(label, band_definition):
-        """Check that a proposed band definition is properly formatted.
+        """Check that a proposed band definition is valid.
 
         Parameters
         ----------
         label : str
-            The name of the new oscillation band.
+            The name of the new band.
         band_definition : tuple of (float, float)
             The lower and upper frequency limit of the band, in Hz.
 
         Raises
         ------
         InconsistentDataError
-            If oscillation band definition is not properly formatted.
+            If band definition is not properly formatted.
         """
 
         # Check that band name is a string

fooof/tests/conftest.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 from fooof.core.modutils import safe_import
-from fooof.tests.utils import get_tfm, get_tfg
+from fooof.tests.utils import get_tfm, get_tfg, get_tbands
 
 plt = safe_import('.pyplot', 'matplotlib')
 
@@ -45,6 +45,10 @@ def tfm():
 def tfg():
     yield get_tfg()
 
+@pytest.fixture(scope='session')
+def tbands():
+    yield get_tbands()
+
 @pytest.fixture(scope='session')
 def skip_if_no_mpl():
     if not safe_import('matplotlib'):

fooof/tests/test_bands.py

@@ -0,0 +1,48 @@
+"""Test functions for FOOOF bands."""
+
+from py.test import raises
+
+from fooof.bands import *
+
+###################################################################################################
+###################################################################################################
+
+def test_bands():
+
+    bands = Bands()
+    assert isinstance(bands, Bands)
+
+def test_bands_add_band():
+
+    bands = Bands()
+    bands.add_band('test', (5, 10))
+    assert bands.bands == {'test' : (5, 10)}
+
+def test_bands_remove_band():
+
+    bands = Bands()
+    bands.add_band('test', (5, 10))
+    bands.remove_band('test')
+    assert bands.bands == {}
+
+def test_bands_errors():
+
+    bands = Bands()
+    with raises(InconsistentDataError):
+        bands.add_band(1, (1, 1))
+    with raises(InconsistentDataError):
+        bands.add_band('test', (1, 1, 1))
+    with raises(InconsistentDataError):
+        bands.add_band('test', (2, 1))
+
+def test_bands_dunders(tbands):
+
+    assert tbands['theta']
+    assert tbands.alpha
+    assert repr(tbands)
+    assert len(tbands) == 3
+
+def test_bands_properties(tbands):
+
+    assert tbands.labels == ['theta', 'alpha', 'beta']
+    assert tbands.n_bands == 3

fooof/tests/utils.py

@@ -3,6 +3,7 @@
 from functools import wraps
 
 from fooof import FOOOF, FOOOFGroup
+from fooof.bands import Bands
 from fooof.synth import gen_power_spectrum, gen_group_power_spectra, param_sampler
 from fooof.core.modutils import safe_import
 
@@ -36,6 +37,11 @@ def get_tfg():
 
     return tfg
 
+def get_tbands():
+    """Get a bands object, for testing."""
+
+    return Bands({'theta' : (4, 8), 'alpha' : (8, 12), 'beta' : (13, 30)})
+
 def default_group_params():
     """Create default parameters for generating a test group of power spectra."""