Repository: mitre/menelaus (language: python)

Issue:
Convert injection functions into callable classes for testing
If we convert these functions into something like:

```python
class ClassJoin():
    def __init__(self):
        # if there really need to be some initial state variables, place here
  
    def __call__(self, data, from, to, c1, c2, new_c):
        self.cls_idx = []
        ...
        return []
```

then we can probe the class' state when the internal calculations pile up, something like:

```python
def test_():
    ...
    assert self.cls_idx = [...]
```

Seems we have the go-ahead from the team on this, so those will be the two efforts in this issue.


Discussion / hints:

Fix patch:
diff --git a/menelaus/injection/__init__.py b/menelaus/injection/__init__.py
index 91127a1f..f550def4 100644
--- a/menelaus/injection/__init__.py
+++ b/menelaus/injection/__init__.py
@@ -1,13 +1,16 @@
-"""This module contains functions with which to artificially inject drift into a dataset."""
+"""
+This module contains classes with which to artificially inject drift into a dataset.
+"""
 
 from menelaus.injection.feature_manipulation import (
-    feature_swap,
-    feature_hide_and_sample,
+    FeatureSwapInjector,
+    FeatureCoverInjector,
 )
-from menelaus.injection.class_manipulation import (
-    class_swap,
-    class_join,
-    class_dirichlet_shift,
-    class_probability_shift,
+from menelaus.injection.injector import Injector
+from menelaus.injection.label_manipulation import (
+    LabelSwapInjector,
+    LabelJoinInjector,
+    LabelProbabilityInjector,
+    LabelDirichletInjector,
 )
-from menelaus.injection.noise import brownian_noise, random_walk
+from menelaus.injection.noise import BrownianNoiseInjector
diff --git a/menelaus/injection/class_manipulation.py b/menelaus/injection/class_manipulation.py
deleted file mode 100644
index 3004d395..00000000
--- a/menelaus/injection/class_manipulation.py
+++ /dev/null
@@ -1,251 +0,0 @@
-import numpy as np
-import pandas as pd
-
-
-# region - simple class manipulation functions
-
-
-def class_swap(data, target_col, from_index, to_index, class_1, class_2):
-    """
-    Swaps two classes in a target column of a given dataset with each other.
-    Accepts ``pandas.DataFrame`` with column names or ``numpy.ndarray``
-    with column indices.
-
-    Ref. :cite:t:`souza2020challenges`
-
-    Args:
-        data (np.array): data to inject with drift
-        target_col (int or str): column index/label of targets column
-        from_index: row index at which to start class swap
-        to_index: row index at which to end (non-inclusive) class swap
-        class_1 (int): value of first label in class swap
-        class_2 (int): value of second label in class swap
-
-    Returns:
-        np.array or pd.DataFrame: copy of data, with two classes swapped
-            in given target column, over given indices
-    """
-    if not isinstance(data, pd.DataFrame) and not isinstance(data, np.ndarray):
-        raise ValueError(f"Data of type {type(data)} not supported")
-    elif isinstance(data, pd.DataFrame):
-        ret = data.copy()
-        class_1_idx = ret.loc[
-            (ret.index < to_index)
-            & (ret.index >= from_index)
-            & (ret[target_col] == class_1)
-        ].index
-        class_2_idx = ret.loc[
-            (ret.index < to_index)
-            & (ret.index >= from_index)
-            & (ret[target_col] == class_2)
-        ].index
-        ret.loc[class_1_idx, target_col] = class_2
-        ret.loc[class_2_idx, target_col] = class_1
-    else:
-        ret = np.copy(data)
-        # np.where returns tuples pertaining to dimensions, only need the row
-        class_1_idx = np.where(ret[:, target_col] == class_1)[0]
-        class_1_idx = class_1_idx[
-            (class_1_idx < to_index) & (class_1_idx >= from_index)
-        ]
-        class_2_idx = np.where(ret[:, target_col] == class_2)[0]
-        class_2_idx = class_2_idx[
-            (class_2_idx < to_index) & (class_2_idx >= from_index)
-        ]
-        ret[class_1_idx, target_col] = class_2
-        ret[class_2_idx, target_col] = class_1
-    return ret
-
-
-def class_join(data, target_col, from_index, to_index, class_1, class_2, new_class):
-    """
-    Joins two [TODO or more?] classes in a unique class. Accepts
-    ``pandas.DataFrame`` with column names or ``numpy.ndarray``
-    with column indices.
-
-    Ref. :cite:t:`souza2020challenges`
-
-    Args:
-        data (np.array): data to inject with drift
-        target_col (int or str): column index/label of targets column
-        from_index: row index at which to start class join
-        to_index: row index at which to end (non-inclusive) class join
-        class_1 (int): value of first label in class join
-        class_2 (int): value of second label in class join,
-        new_class (int): new label value to assign to old classes
-
-    Returns:
-        np.array or pd.DataFrame: copy of data, with two classes joined
-            in given target column, over given indices, into new class
-    """
-    if not isinstance(data, pd.DataFrame) and not isinstance(data, np.ndarray):
-        raise ValueError(f"Data of type {type(data)} not supported")
-    elif isinstance(data, pd.DataFrame):
-        ret = data.copy()
-        class_idx = ret.loc[
-            (ret.index < to_index)
-            & (ret.index >= from_index)
-            & ((ret[target_col] == class_1) | (ret[target_col] == class_2))
-        ].index
-        ret.loc[class_idx, target_col] = new_class
-    else:
-        ret = np.copy(data)
-        # np.where returns tuples pertaining to dimensions, only need the row
-        class_idx = np.where(
-            (ret[:, target_col] == class_1) | (ret[:, target_col] == class_2)
-        )[0]
-        class_idx = class_idx[(class_idx < to_index) & (class_idx >= from_index)]
-
-        ret[class_idx, target_col] = new_class
-    return ret
-
-
-# endregion
-
-
-# region - LTF-inspired class manipulation functions
-
-
-def class_probability_shift(
-    data, target_col, from_index, to_index, class_probabilities
-):
-    """
-    Resamples the data over a specified window, with altered probability
-    for specified classes (and uniform probability for remaining classes).
-    Accepts ``pandas.DataFrame`` with column names or ``numpy.ndarray``
-    with column indices.
-
-    Note:
-    * this function can perform tweak-one and minority shift
-    * When a class is not present in the window specified,
-    but specified in ``class_probabilities``, the probability
-    value is uniformly divided into the remaining classes in
-    the window.
-
-    Ref. :cite:t:`LTFmethods`
-
-    Args:
-        data (np.array): data to inject with drift
-        target_col (int or str): column index/label of targets column
-        from_index (int): row index at which to start class swap
-        to_index (int): row index at which to end (non-inclusive) class swap
-        class_probabilities (dict): classes as keys, and their desired
-            resampling chance as values. Un-specified classes are given
-            a uniform resampling chance with respect to all other
-            un-specified classes
-
-    Returns:
-        np.array or pd.DataFrame: copy of data, resampled with shifted
-            class probability for 1 or more desired classes
-    """
-    # if pandas.DataFrame, store columns, convert to numpy
-    columns = None
-    if isinstance(data, pd.DataFrame):
-        columns = data.columns
-        data = data.to_numpy()
-        ret = data.copy()
-    elif isinstance(data, np.ndarray):
-        ret = np.copy(data)
-    else:
-        raise ValueError(f"Data of type {type(data)} not supported")
-
-    # determine all unique classes and classes not specified in args
-    all_classes = np.unique(ret[:, target_col])
-    undefined_classes = [k for k in all_classes if k not in class_probabilities]
-
-    # specified class probabilities must sum to 1 or less
-    if sum(class_probabilities.values()) > 1.0:
-        raise ValueError(f"Probabilities in {class_probabilities} exceed 1")
-
-    # args should not specify previously unseen classes
-    if set(all_classes) != set(list(class_probabilities.keys()) + undefined_classes):
-        raise ValueError(
-            f"Argument {class_probabilities} has classes not found in data {all_classes}"
-        )
-
-    # undefined classes are resampled uniformly
-    for uc in undefined_classes:
-        class_probabilities[uc] = (1 - sum(class_probabilities.values())) / len(
-            undefined_classes
-        )
-
-    # distribution for each data point, and reordering of each point by class
-    p_distribution = []
-    sample_idxs_grouped = []
-
-    # locate each class in window
-    for cls in all_classes:
-        cls_idx = np.where(data[:, target_col] == cls)[0]
-        cls_idx = cls_idx[(cls_idx < to_index) & (cls_idx >= from_index)]
-
-        # each member has p_class / class_size chance, represented as bool to avoid div/0
-        p_individual = (
-            cls_idx.shape[0] and class_probabilities[cls] / cls_idx.shape[0]
-        ) or 0
-
-        # append to grouped array and corresponding distribution
-        sample_idxs_grouped.extend(cls_idx)
-        p_distribution.extend(np.ones(cls_idx.shape[0]) * p_individual)
-
-    # if classes skipped, ensure probability distribution adds to 1
-    p_leftover = (1 - sum(p_distribution)) / len(p_distribution)
-    p_distribution = [p + p_leftover for p in p_distribution]
-
-    # shuffled sample over window, with replacement, with weights
-    sample_idxs = np.random.choice(
-        sample_idxs_grouped, to_index - from_index, True, p_distribution
-    )
-    ret[from_index:to_index] = data[sample_idxs]
-
-    # back to DF if needed
-    if columns is not None:
-        ret = pd.DataFrame(ret)
-        ret.columns = columns
-
-    return ret
-
-
-def class_dirichlet_shift(data, target_col, from_index, to_index, alpha):
-    """
-    Resamples the data over a specified window, per a generated Dirichlet
-    distribution (with specified alpha) over all labels.
-    Accepts ``pandas.DataFrame`` with column names or ``numpy.ndarray``
-    with column indices.
-
-    Note: If all labels are not given weights, unexpected behavior may
-    cause all un-specified classes to be given uniform resampling chance.
-
-    Ref. :cite:t:`LTFmethods`
-
-    Args:
-        data (np.array): data to inject with drift
-        target_col (int or str): column index/label of targets column
-        from_index: row index at which to start class swap
-        to_index: row index at which to end (non-inclusive) class swap
-        alpha (dict): used to derive alpha parameter for Dirichlet
-            distribution. Keys are ALL labels, values are the desired
-            average weight (typically ``int``) per label when resampling.
-            For example, weights of [4,1] correspond to an expected 80/20
-            percent split between first and second classes.
-
-    Returns:
-        np.array or pd.DataFrame: copy of data, resampled per Dirichlet
-            distribution over classes with specified alpha
-    """
-    alpha_classes = list(alpha.keys())
-    alpha_values = [alpha[k] for k in alpha]
-
-    # generate dirichlet distribution by class
-    # XXX - minor concern that order of these list-types not always guaranteed
-    dirichlet_distribution = np.random.dirichlet(alpha_values)
-    dirichlet_probabilities = {
-        alpha_classes[i]: dirichlet_distribution[i] for i in range(len(alpha_classes))
-    }
-
-    # use class_probability_shift with fully-specified distribution
-    return class_probability_shift(
-        data, target_col, from_index, to_index, dirichlet_probabilities
-    )
-
-
-# endregion
diff --git a/menelaus/injection/feature_manipulation.py b/menelaus/injection/feature_manipulation.py
index 97e61d60..9da8e6df 100644
--- a/menelaus/injection/feature_manipulation.py
+++ b/menelaus/injection/feature_manipulation.py
@@ -1,68 +1,74 @@
-import numpy as np
-import pandas as pd
+from menelaus.injection.injector import Injector
 
 
-def feature_swap(data, col_1, col_2, from_index, to_index):
+class FeatureSwapInjector(Injector):
     """
     Swaps two features/columns of a given dataset with each other.
     Accepts ``pandas.DataFrame`` with column names or ``numpy.ndarray``
     with column indices.
 
     Ref. :cite:t:`souza2020challenges`
+    """
 
-    Args:
-        data (np.ndarray or pd.DataFrame): data to inject with drift
-        col_1 (int or str): column index/label of first column
-        col_2 (int or str): column index/label of second column
-        from_index (int): row index at which to start column swap
-        to_index (int): row index at which to end (non-inclusive) column swap
+    def __call__(self, data, from_index, to_index, col_1, col_2):
+        """
+        Args:
+            data (np.ndarray or pd.DataFrame): data to inject with drift
+            from_index (int): row index at which to start column swap
+            to_index (int): row index at which to end (non-inclusive) column swap
+            col_1 (int or str): column index/name of first column
+            col_2 (int or str): column index/name of second column
 
-    Returns:
-        np.ndarray or pd.DataFrame: copy of data, with two columns swapped
-            over given indices
-    """
-    if not isinstance(data, pd.DataFrame) and not isinstance(data, np.ndarray):
-        raise ValueError(f"Data of type {type(data)} not supported")
-    elif isinstance(data, pd.DataFrame):
-        ret = data.copy()
-        ret_col_1 = ret.loc[from_index:to_index, col_1].copy()
-        ret.loc[from_index:to_index, col_1] = ret.loc[from_index:to_index, col_2]
-        ret.loc[from_index:to_index, col_2] = ret_col_1
-    else:
-        ret = np.copy(data)
+        Returns:
+            np.ndarray or pd.DataFrame: copy of data, with two columns swapped
+                over given indices
+        """
+        # handle type
+        ret, (col_1, col_2) = self._preprocess(data, col_1, col_2)
+
+        # swap columns
         ret[from_index:to_index, [col_1, col_2]] = ret[
             from_index:to_index, [col_2, col_1]
         ]
-    return ret
 
+        # handle type and return
+        ret = self._postprocess(ret)
+        return ret
 
-def feature_hide_and_sample(data, col, sample_size, random_state=0):
+
+class FeatureCoverInjector(Injector):
     """
     Hides a feature, then treats it as a shared concept by which to group the data.
     Afterwards samples are uniformly drawn from each group. Accepts
     ``pandas.DataFrame`` with column names or ``numpy.ndarray`` with column indices.
 
+    Note:
+    * This function cannot hide a feature over a window, as the appearance of the
+        column outside the window may raise errors.
+
     Ref. :cite:t:`souza2020challenges`
+    """
 
-    Args:
-        data (np.ndarray or pd.DataFrame): data to inject with drift
-        col (int or str): index/label of column to hide and re-sample (note this
-            should be a categorical feature that can be treated as a concept)
-        sample_size (int): data points to be drawn from each group in new concept
-        random_state (int): optional random seed. Default 0
+    def __call__(self, data, col, sample_size, random_state=None):
+        """
+        Args:
+            data (np.ndarray or pd.DataFrame): data to inject with drift
+            col (int or str): index/label of column to hide and re-sample (note this
+                should be a categorical feature that can be treated as a concept)
+            sample_size (int): data points to be drawn from each group in new concept
+            random_state (int): optional random seed. Default 0
 
-    Returns:
-        np.ndarray or pd.DataFrame: copy of data, grouped by indicated column,
-            with each group sampled and column removed
-    """
-    if not isinstance(data, pd.DataFrame) and not isinstance(data, np.ndarray):
-        raise ValueError(f"Data of type {type(data)} not supported")
-    elif isinstance(data, pd.DataFrame):
-        df = data.copy()
-    else:
-        df = np.copy(data)
-        df = pd.DataFrame(df)
-    ret = df.groupby(col).sample(n=sample_size, random_state=random_state)
-    ret = ret.drop(columns=[col]).reset_index(drop=True)
-    ret = ret if isinstance(data, pd.DataFrame) else ret.to_numpy()
-    return ret
+        Returns:
+            np.ndarray or pd.DataFrame: copy of data, grouped by indicated column,
+                with each group sampled and column removed
+        """
+        # handle type
+        ret, (col,) = self._preprocess(data, col, return_df=True)
+
+        # hide and reorder
+        ret = ret.groupby(col).sample(n=sample_size, random_state=random_state)
+        ret = ret.drop(columns=[col]).reset_index(drop=True)
+
+        # handle type and return
+        ret = self._postprocess(ret)
+        return ret
diff --git a/menelaus/injection/injector.py b/menelaus/injection/injector.py
new file mode 100644
index 00000000..600a075f
--- /dev/null
+++ b/menelaus/injection/injector.py
@@ -0,0 +1,99 @@
+from abc import ABC, abstractclassmethod
+
+import numpy as np
+import pandas as pd
+
+
+class Injector(ABC):
+    """
+    Abstract base class for drift injection callables. Classes using this
+    pattern should implement a ``__call__`` function, and also make use of the
+    ``_preprocess()`` and ``_postprocess`` type handling capabilities of this
+    class.
+
+    When called, an ``Injector`` typically accepts either ``pandas.DataFrame``
+    or ``numpy.ndarray`` data, along with other parameters for the injection
+    technique it is modeling, and returns the data with drift injected over a
+    specified window.
+    """
+
+    def _preprocess(self, data, *columns, return_df=False):
+        """
+        Preprocesses data and any column indicators into desired format. By
+        default returns data as ``numpy.ndarray`` and any columns as integer
+        column indices unless told otherwise.
+
+        Note:
+        * ``*columns`` is used to preprocess any columns used by a sub-class'
+        ``__call__`` function, and does not necessarily represent ALL columns
+        * if the input is in ``numpy`` format, and a ``numpy.ndarray`` is
+        desired, ``*columns`` should be integers
+        * ``*columns`` should only be string names if the input is a
+        ``pandas.DataFrame``
+
+        Args:
+            data (numpy.ndarray or pd.DataFrame): data to preprocess
+            *columns (str or int): 1 or more integers or names corresponding to
+                columns in the data, to also be preprocessed
+            return_df (bool): Whether to return data as ``pandas.DataFrame``,
+                in which case ``*columns`` are returned as-is. Default False.
+
+        Returns:
+            tuple: first item is a deep-copy of the data in ``numpy.ndarray``,
+                unless ``return_df`` is True, in which case a
+                ``pandas.DataFrame`` is returned. Second item is the column
+                indicators as integer indexes, unless ``return_df`` is True, in
+                which case the column indicators should be strings and are
+                returned as-is.
+        """
+        # store columns if DF, otherwise set to None
+        if isinstance(data, np.ndarray):
+            self._columns = None
+            column_idxs = columns
+
+        # column str names to integer idxs if needed
+        elif isinstance(data, pd.DataFrame):
+            self._columns = data.columns
+            column_idxs = tuple([data.columns.get_loc(c) for c in columns])
+
+        # only two types supported
+        else:
+            raise ValueError(f"Data of type {type(data)} not supported")
+
+        # copy and return desired type
+        copy = np.copy(data)
+        if return_df:
+            return pd.DataFrame(copy, columns=self._columns), columns
+        else:
+            return copy, column_idxs
+
+    def _postprocess(self, data):
+        """
+        Postprocesses data and returns as its initial type, which was
+        determined and recorded during preprocessing.
+
+        Args:
+            data (``numpy.ndarray`` or ``pandas.DataFrame``): data to
+                postprocess
+
+        Returns:
+            data (``numpy.ndarray`` ``or pandas.DataFrame``): data,
+                transformed into its original state
+        """
+        # if column names were stored and data is not already DF, return DF
+        if self._columns is not None and not isinstance(data, pd.DataFrame):
+            return pd.DataFrame(data, columns=self._columns)
+
+        # if column names not stored and data is DF, return original ndarray
+        elif self._columns is None and isinstance(data, pd.DataFrame):
+            return data.to_numpy()
+
+        # if column names stored and data is DF, return unchanged
+        # if column names not stored and data is ndarray, return unchanged
+        else:
+            return data
+
+    @abstractclassmethod
+    def __call__(self):
+        """Implemented by sub-classes."""
+        raise NotImplementedError
diff --git a/menelaus/injection/label_manipulation.py b/menelaus/injection/label_manipulation.py
new file mode 100644
index 00000000..1de5e0fd
--- /dev/null
+++ b/menelaus/injection/label_manipulation.py
@@ -0,0 +1,242 @@
+import numpy as np
+
+from menelaus.injection.injector import Injector
+
+
+# region - Simple Label Manipulation
+
+
+class LabelSwapInjector(Injector):
+    """
+    Swaps two classes in a target column of a given dataset with each other.
+    Accepts ``pandas.DataFrame`` with column names or ``numpy.ndarray``
+    with column indices.
+
+    Ref. :cite:t:`souza2020challenges`
+    """
+
+    def __call__(self, data, from_index, to_index, target_col, class_1, class_2):
+        """
+        Args:
+            data (np.array): data to inject with drift
+            from_index: row index at which to start class swap
+            to_index: row index at which to end (non-inclusive) class swap
+            target_col (int or str): column index/name of targets column
+            class_1 (int): value of first label in class swap
+            class_2 (int): value of second label in class swap
+
+        Returns:
+            np.array or pd.DataFrame: copy of data, with two classes swapped
+                in given target column, over given indices
+        """
+        # handle data type
+        ret, (target_col,) = self._preprocess(data, target_col)
+
+        # locate two classes, perform swap
+        class_1_idx = np.where(ret[:, target_col] == class_1)[0]
+        class_1_idx = class_1_idx[
+            (class_1_idx < to_index) & (class_1_idx >= from_index)
+        ]
+        class_2_idx = np.where(ret[:, target_col] == class_2)[0]
+        class_2_idx = class_2_idx[
+            (class_2_idx < to_index) & (class_2_idx >= from_index)
+        ]
+        ret[class_1_idx, target_col] = class_2
+        ret[class_2_idx, target_col] = class_1
+
+        # handle data type and return
+        ret = self._postprocess(ret)
+        return ret
+
+
+class LabelJoinInjector(Injector):
+    """
+    Joins two [TODO or more?] classes in a unique class. Accepts
+    ``pandas.DataFrame`` with column names or ``numpy.ndarray``
+    with column indices.
+
+    Ref. :cite:t:`souza2020challenges`
+    """
+
+    def __call__(
+        self, data, from_index, to_index, target_col, class_1, class_2, new_class
+    ):
+        """
+        Args:
+            data (np.array): data to inject with drift
+            from_index: row index at which to start class join
+            to_index: row index at which to end (non-inclusive) class join
+            target_col (int or str): column index/name of targets column
+            class_1 (int): value of first label in class join
+            class_2 (int): value of second label in class join,
+            new_class (int): new label value to assign to old classes
+
+        Returns:
+            np.array or pd.DataFrame: copy of data, with two classes joined
+                in given target column, over given indices, into new class
+        """
+        # handle data type
+        ret, (target_col,) = self._preprocess(data, target_col)
+
+        # locate two labels, switch both to new label
+        class_idx = np.where(
+            (ret[:, target_col] == class_1) | (ret[:, target_col] == class_2)
+        )[0]
+        class_idx = class_idx[(class_idx < to_index) & (class_idx >= from_index)]
+        ret[class_idx, target_col] = new_class
+
+        # handle data type and return
+        ret = self._postprocess(ret)
+        return ret
+
+
+# endregion
+
+# region - LTF-Inspired Label Manipulation
+
+
+class LabelProbabilityInjector(Injector):
+    """
+    Resamples the data over a specified window, with altered probability
+    for specified classes (and uniform probability for remaining classes).
+    Accepts ``pandas.DataFrame`` with column names or ``numpy.ndarray``
+    with column indices.
+
+    Note:
+    * this function can perform tweak-one and minority shift
+    * When a class is not present in the window specified,
+    but specified in ``class_probabilities``, the probability
+    value is uniformly divided into the remaining classes in
+    the window.
+
+    Ref. :cite:t:`LTFmethods`
+    """
+
+    def __call__(self, data, from_index, to_index, target_col, class_probabilities):
+        """
+        Args:
+            data (np.array): data to inject with drift
+            from_index (int): row index at which to start class swap
+            to_index (int): row index at which to end (non-inclusive) class swap
+            target_col (int or str): column index/name of targets column
+            class_probabilities (dict): classes as keys, and their desired
+                resampling chance as values. Un-specified classes are given
+                a uniform resampling chance with respect to all other
+                un-specified classes
+
+        Returns:
+            np.array or pd.DataFrame: copy of data, resampled with shifted
+                class probability for 1 or more desired classes
+        """
+        # handle data type
+        ret, (target_col,) = self._preprocess(data, target_col)
+
+        # determine all unique classes and classes not specified in args
+        all_classes = np.unique(ret[:, target_col])
+        undefined_classes = [k for k in all_classes if k not in class_probabilities]
+
+        # specified class probabilities must sum to 1 or less
+        if sum(class_probabilities.values()) > 1.0:
+            raise ValueError(f"Probabilities in {class_probabilities} exceed 1")
+
+        # args should not specify previously unseen classes
+        if set(all_classes) != set(
+            list(class_probabilities.keys()) + undefined_classes
+        ):
+            raise ValueError(
+                f"Argument {class_probabilities} has classes not found in data {all_classes}"
+            )
+
+        # undefined classes are resampled uniformly
+        missing_probability = 1 - sum(class_probabilities.values())
+        for uc in undefined_classes:
+            class_probabilities[uc] = missing_probability / len(undefined_classes)
+
+        # distribution for each data point, and reordering of each point by class
+        self._p_distribution = []
+        sample_idxs_grouped = []
+
+        # locate each class in window
+        for cls in all_classes:
+            cls_idx = np.where(data[:, target_col] == cls)[0]
+            cls_idx = cls_idx[(cls_idx < to_index) & (cls_idx >= from_index)]
+
+            # each member has p_class / class_size chance, represented as bool to avoid div/0
+            p_individual = (
+                cls_idx.shape[0] and class_probabilities[cls] / cls_idx.shape[0]
+            ) or 0
+
+            # append to grouped array and corresponding distribution
+            sample_idxs_grouped.extend(cls_idx)
+            self._p_distribution.extend(np.ones(cls_idx.shape[0]) * p_individual)
+
+        # if classes skipped, ensure probability distribution adds to 1
+        p_leftover = (1 - sum(self._p_distribution)) / len(self._p_distribution)
+        self._p_distribution = [p + p_leftover for p in self._p_distribution]
+
+        # shuffled sample over window, with replacement, with weights
+        sample_idxs = np.random.choice(
+            sample_idxs_grouped, to_index - from_index, True, self._p_distribution
+        )
+        ret[from_index:to_index] = data[sample_idxs]
+
+        # handle data type and return
+        ret = self._postprocess(ret)
+        return ret
+
+
+class LabelDirichletInjector(Injector):
+    """
+    Resamples the data over a specified window, per a generated Dirichlet
+    distribution (with specified alpha) over all labels. Accepts
+    ``pandas.DataFrame`` with column names or ``numpy.ndarray``
+    with column indices.
+
+    Notes:
+    * If all labels are not given weights, unexpected behavior may
+    cause all un-specified classes to be given uniform resampling chance.
+
+    Ref. :cite:t:`LTFmethods`
+    """
+
+    def __call__(self, data, from_index, to_index, target_col, alpha):
+        """
+        Args:
+            data (np.array): data to inject with drift
+            from_index: row index at which to start class swap
+            to_index: row index at which to end (non-inclusive) class swap
+            target_col (int or str): column index/name of targets column
+            alpha (dict): used to derive alpha parameter for Dirichlet
+                distribution. Keys are ALL labels, values are the desired
+                average weight (typically ``int``) per label when resampling.
+                For example, weights of [4,1] correspond to an expected 80/20
+                percent split between first and second classes.
+
+        Returns:
+            np.array or pd.DataFrame: copy of data, resampled per Dirichlet
+                distribution over classes with specified alpha
+        """
+        # weights dictionary into lists of keys and values
+        self._alpha_classes = list(alpha.keys())
+        self._alpha_values = [alpha[k] for k in alpha]
+
+        # generate dirichlet distribution by class
+        # XXX - minor concern that order of these list-types not always guaranteed
+        self._dirichlet_distribution = np.random.dirichlet(self._alpha_values)
+        self._dirichlet_probabilities = {
+            self._alpha_classes[i]: self._dirichlet_distribution[i]
+            for i in range(len(self._alpha_classes))
+        }
+
+        # use class_probability_shift with fully-specified distribution
+        label_prob_injector = LabelProbabilityInjector()
+        return label_prob_injector(
+            data,
+            from_index=from_index,
+            to_index=to_index,
+            target_col=target_col,
+            class_probabilities=self._dirichlet_probabilities,
+        )
+
+
+# endregion
diff --git a/menelaus/injection/noise.py b/menelaus/injection/noise.py
index be512461..c2e21295 100644
--- a/menelaus/injection/noise.py
+++ b/menelaus/injection/noise.py
@@ -1,55 +1,62 @@
 import numpy as np
-import pandas as pd
 
+from menelaus.injection.injector import Injector
 
-def brownian_noise(data, feature, x0, from_index, to_index):
+
+class BrownianNoiseInjector(Injector):
     """
     Adds Brownian noise to a column of a given dataset. Accepts
     ``pandas.DataFrame`` with column names or ``numpy.ndarray``
     with column indices. Calls ``random_walk`` for full computation.
 
     Ref. :cite:t:`fields2019mitigating`
-
-    Args:
-        data (np.ndarray or pd.DataFrame): data to inject with drift
-        feature (int or str): index/label of column to add noise to
-        x0 (int): starting point for random walk noise
-        from_index (int): row index at which to start injection
-        to_index (int): row index at which to end (non-inclusive) injection
-
-    Returns:
-        np.ndarray or pd.DataFrame: copy of data, with two columns swapped
-            over given indices
-    """
-    steps = to_index - from_index
-    if not isinstance(data, pd.DataFrame) and not isinstance(data, np.ndarray):
-        raise ValueError(f"Data of type {type(data)} not supported")
-    elif isinstance(data, pd.DataFrame):
-        ret = data.copy()
-        # unlike other operations, .loc[0:r] is inclusive, so -1
-        ret_col = ret.loc[from_index : to_index - 1, feature].copy()
-        ret_col = ret_col + random_walk(steps, x0)
-        ret.loc[from_index:to_index, feature] = ret_col
-    else:
-        ret = np.copy(data)
-        ret[from_index:to_index, feature] = ret[
-            from_index:to_index, feature
-        ] + random_walk(steps, x0)
-    return ret
-
-
-def random_walk(steps, x0):
     """
-    Performs Weiner process given a number of steps and initial starting point.
 
-    Ref. :cite:t:`fields2019mitigating`
-
-    Args:
-        steps (int): number of steps, corresponds to size of output vector
-        x0 (int): initial starting value for output vector values
-    """
-    w = np.ones(steps) * x0
-    for i in range(1, steps):
-        yi = np.random.choice([1, -1])
-        w[i] = w[i - 1] + (yi / np.sqrt(steps))
-    return w
+    def __call__(self, data, from_index, to_index, col, x0, random_state=None):
+        """
+        Args:
+            data (np.ndarray or pd.DataFrame): data to inject with drift
+            from_index (int): row index at which to start injection
+            to_index (int): row index at which to end (non-inclusive) injection
+            col (int or str): index/name of column to add noise to
+            x0 (int): starting point for random walk noise
+            random_state (int): optional value for setting seed. Default None.
+
+        Returns:
+            np.ndarray or pd.DataFrame: copy of data, with Brownian noise
+                injected over given indices
+        """
+        # handle type
+        ret, (col,) = self._preprocess(data, col)
+
+        # add RW noise
+        steps = to_index - from_index
+        ret[from_index:to_index, col] = ret[
+            from_index:to_index, col
+        ] + BrownianNoiseInjector._random_walk(steps, x0, random_state=random_state)
+
+        # handle type and return
+        ret = self._postprocess(ret)
+        return ret
+
+    @staticmethod
+    def _random_walk(steps, x0, random_state=None):
+        """
+        Performs Weiner process given a number of steps and initial starting point.
+
+        Ref. :cite:t:`fields2019mitigating`
+
+        Args:
+            steps (int): number of steps, corresponds to size of output vector
+            x0 (int): initial starting value for output vector values
+            random_state (int): optional value for setting seed. Default None.
+        """
+        # set seed if given
+        np.random.seed(random_state)
+
+        # perform random walk
+        w = np.ones(steps) * x0
+        for i in range(1, steps):
+            yi = np.random.choice([1, -1])
+            w[i] = w[i - 1] + (yi / np.sqrt(steps))
+        return w


Test patch:
diff --git a/tests/menelaus/injection/test_class_manipulation.py b/tests/menelaus/injection/test_class_manipulation.py
deleted file mode 100644
index 06f4c422..00000000
--- a/tests/menelaus/injection/test_class_manipulation.py
+++ /dev/null
@@ -1,131 +0,0 @@
-import pytest
-import pandas as pd
-import numpy as np
-
-
-from menelaus.injection import class_join, class_swap, class_probability_shift, class_dirichlet_shift
-
-
-def test_join_1():
-    ''' Check correct join in pd.DataFrame '''
-    data = pd.DataFrame({'a': [0,0,1,1], 'b': [5,5,5,5]})
-    expected_data = pd.DataFrame({'a': [0,2,2,1], 'b':[5,5,5,5]})
-    data = class_join(data, 'a', 1, 3, 0, 1, 2)
-    assert data.equals(expected_data)
-
-def test_join_2():
-    ''' Check correct join in numpy.ndarray '''
-    data = np.array([[0,5], [0,5], [1,5], [1,5]])
-    expected_data = np.array([[0,5], [2,5], [2,5], [1,5]])
-    data = class_join(data, 0, 1, 3, 0, 1, 2)
-    assert np.array_equal(data, expected_data)
-
-def test_join_3():
-    ''' Check ValueError if data neither pandas.DataFrame nor np.ndarray '''
-    data = [[0,0], [1,1], [0,2], [1,3]]
-    with pytest.raises(ValueError):
-        class_join(
-            data, target_col=0, class_1=0, class_2=1,
-            new_class=2, from_index=1, to_index=3
-        )
-
-def test_swap_1():
-    ''' Check correct swap in pd.DataFrame '''
-    data = pd.DataFrame({'a': [0,0,1,1], 'b': [5,5,6,6]})
-    expected_data = pd.DataFrame({'a': [1,1,0,0], 'b':[5,5,6,6]})
-    data = class_swap(data, 'a', 0, 4, 0, 1)
-    assert data.equals(expected_data)
-
-def test_swap_2():
-    ''' Check correct swap in numpy.ndarray'''
-    data = np.array([[0,5], [0,5], [1,6], [1,6]])
-    expected_data = np.array([[1,5], [1,5], [0,6], [0,6]])
-    data = class_swap(data, 0, 0, 4, 0, 1)
-    assert np.array_equal(data, expected_data)
-
-def test_swap_3():
-    ''' Check ValueError if data neither pandas.DataFrame nor np.ndarray '''
-    data = [[0,0], [1,1], [0,2], [1,3]]
-    with pytest.raises(ValueError):
-        class_swap(
-            data, target_col=0, class_1=0, class_2=1, 
-            from_index=1, to_index=3
-        )
-
-def test_probability_shift_1():
-    ''' Ensure probability shift causes some drift in numpy.ndarray '''
-    np.random.seed(0)
-    data = np.array([
-        [0,0,0], [1,0,0], [2,0,0], [3,0,0], [4,0,0], [5,0,0],
-        [6,0,0], [7,0,0], [8,0,0], [9,0,0]
-    ])
-    probs = {k:0.1 for k in range(0,10)}
-    new_data = class_probability_shift(data, 0, 1, 9, probs)
-    assert np.array_equal(data[0], new_data[0])
-    assert np.array_equal(data[-1], new_data[-1])
-    assert not np.array_equal(data[1:9], new_data[1:9])
-
-def test_probability_shift_2():
-    ''' Ensure probability shift causes some drift in pandas.DataFrame '''
-    np.random.seed(0)
-    data = pd.DataFrame({
-        'a': [0,1,2,3,4,5,6,7,8,9],
-        'b': [0,0,0,0,0,0,0,0,0,0]
-    })
-    probs = {k:0.1 for k in range(0,10)}
-    new_data = class_probability_shift(data, 0, 1, 9, probs)
-    assert data.iloc[0].equals(new_data.iloc[0])
-    assert data.iloc[-1].equals(new_data.iloc[-1])
-    assert not np.array_equal(data[1:9], new_data[1:9])
-    assert list(new_data.columns) == ['a', 'b']
-
-def test_probability_shift_3():
-    ''' Check ValueError if data neither pandas.DataFrame nor numpy.ndarray '''
-    data = [[0,0], [1,1], [0,2], [1,3]]
-    with pytest.raises(ValueError):
-        class_probability_shift(
-            data, target_col=0, from_index=0, to_index=2, 
-            class_probabilities={0: 0.5}
-        )
-
-def test_probability_shift_4():
-    ''' Check ValueError when probabilities add to >1 '''
-    data = np.random.randint(0,5,(10,3))
-    with pytest.raises(ValueError):
-        class_1 = np.unique(data[:,0])[0]
-        class_2 = np.unique(data[:,0])[1]
-        probs = {class_1: 1.0, class_2: 0.5}
-        class_probability_shift(data, 0, 0, 3, probs)
-
-def test_probability_shift_5():
-    ''' Check ValueError when given probabilities for classes not in data '''
-    data = np.random.randint(0,5,(10,3))
-    with pytest.raises(ValueError):
-        probs = {1000: 0.5}
-        class_probability_shift(data, 0, 0, 3, probs)
-
-def test_probability_shift_6():
-    ''' Check drift injected when not all classes specified '''
-    np.random.seed(0)
-    data = np.array([
-        [0,0,0], [1,0,0], [2,0,0], [3,0,0], [4,0,0], [5,0,0],
-        [6,0,0], [7,0,0], [8,0,0], [9,0,0]
-    ])
-    probs = {k:0.1 for k in range(2,10)}
-    new_data = class_probability_shift(data, 0, 1, 9, probs)
-    assert np.array_equal(data[0], new_data[0])
-    assert np.array_equal(data[-1], new_data[-1])
-    assert not np.array_equal(data[1:9], new_data[1:9])
-
-def test_dirichlet_shift_1():
-    ''' Ensure Dirichlet shift causes some drift '''
-    np.random.seed(0)
-    data = np.array([
-        [0,0,0], [1,0,0], [2,0,0], [3,0,0], [4,0,0], [5,0,0],
-        [6,0,0], [7,0,0], [8,0,0], [9,0,0]
-    ])
-    alpha = {k:1 for k in range(0,10)}
-    new_data = class_dirichlet_shift(data, 0, 1, 9, alpha)
-    assert np.array_equal(data[0], new_data[0])
-    assert np.array_equal(data[-1], new_data[-1])
-    assert not np.array_equal(data[1:9], new_data[1:9])
\ No newline at end of file
diff --git a/tests/menelaus/injection/test_feature_manipulation.py b/tests/menelaus/injection/test_feature_manipulation.py
index f90a762d..a4df9df5 100644
--- a/tests/menelaus/injection/test_feature_manipulation.py
+++ b/tests/menelaus/injection/test_feature_manipulation.py
@@ -1,49 +1,21 @@
-import pytest
-import pandas as pd
 import numpy as np
 
-from menelaus.injection import feature_swap, feature_hide_and_sample
+from menelaus.injection import FeatureSwapInjector, FeatureCoverInjector
 
 
 def test_feature_swap_1():
-    ''' Check correct swap in pandas.DataFrame '''
-    data = pd.DataFrame({'a':[0,1,2,3,4], 'b':[0,8,9,3,4]})
-    expected_data = pd.DataFrame({'a':[0,8,9,3,4], 'b':[0,1,2,3,4]})
-    data = feature_swap(data, 'a', 'b', 1, 3)
-    assert data.equals(expected_data)
-
-def test_feature_swap_2():
-    ''' Check correct swap in numpy.ndarray '''
+    ''' Check correct feature swap behavior '''
     data = np.array([[0,0], [1,8], [2,9], [3,3], [4,4]])
     expected_data = np.array([[0,0], [8,1], [9,2], [3,3], [4,4]])
-    data = feature_swap(data, 0, 1, 1, 3)
-    assert np.array_equal(data, expected_data)
-
-def test_feature_swap_3():
-    ''' Check function fails when data is neither numpy nor pandas '''
-    data = [1,2,3,4]
-    with pytest.raises(ValueError):
-        feature_swap(data, 'col1', 'col2', 0, 2)
-
-def test_feature_hide_1():
-    ''' Check correct operation in pandas.DataFrame '''
-    data = pd.DataFrame({'a': [0, 0, 0, 1, 1, 1], 'b': [0, 1, 2, 3, 4, 5]})
-    expected_data = pd.DataFrame({'b': [2,1,5,3]}) 
-    # expect result, since random state is default 0 in function call
-    data = feature_hide_and_sample(data, 'a', 2)
-    assert data.equals(expected_data)
-
-def test_feature_hide_2():
-    ''' Check correct operation in numpy.ndarray '''
-    data = np.array([[0,0], [0,1], [0,2], [1,3], [1,4], [1,5]])
-    expected_data = np.array([2,1,5,3]) 
-    # expect result, since random state is default 0 in function call
-    data = feature_hide_and_sample(data, 0, 2)
-    data = data.reshape(-1)
-    assert np.array_equal(data, expected_data)
-
-def test_feature_hide_3():
-    ''' Check function fails when data is neither numpy nor pandas '''
-    data = [1,2,3,4]
-    with pytest.raises(ValueError):
-        feature_hide_and_sample(data, 'col', 2)
\ No newline at end of file
+    i = FeatureSwapInjector()
+    copy = i(data, from_index=1, to_index=3, col_1=0, col_2=1)
+    assert np.array_equal(copy, expected_data)
+
+def test_feature_cover_1():
+    ''' Check correct feature cover behavior '''
+    data = np.array([[0,2], [0,2], [0,2], [1,3], [1,3], [1,3]])
+    i = FeatureCoverInjector()
+    copy = i(data, col=0, sample_size=2, random_state=0)
+    assert copy.shape == (4,1)
+    assert np.array_equal(np.where(copy[:,0]==2)[0], [0,1])
+    assert np.array_equal(np.where(copy[:,0]==3)[0], [2,3])
diff --git a/tests/menelaus/injection/test_injector.py b/tests/menelaus/injection/test_injector.py
new file mode 100644
index 00000000..f04315d4
--- /dev/null
+++ b/tests/menelaus/injection/test_injector.py
@@ -0,0 +1,86 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from menelaus.injection import Injector, LabelSwapInjector
+
+
+# override this attribute so that it is instantiable
+# hacky, but lets us reach 100% coverage
+Injector.__abstractmethods__ = set()
+
+
+def test_preprocess_1():
+    ''' Check behavior when ndarray input, ndarray desired output, column processing needed '''
+    i = LabelSwapInjector()
+    data = np.random.randint(0,5,(4,4))
+    copy, cols = i._preprocess(data, 0, 1)
+    assert np.array_equal(data, copy)
+    assert cols == (0,1)
+
+def test_preprocess_2():
+    ''' Check behavior when ndarray input, dataframe desired output, column processing needed '''
+    i = LabelSwapInjector()
+    data = np.ones((2,2))
+    copy, cols = i._preprocess(data, 0, 1, return_df=True)
+    expected_copy = pd.DataFrame({0: [1.0,1.0], 1: [1.0,1.0]})
+    assert expected_copy.equals(copy)
+    assert cols == (0,1)
+
+def test_preprocess_3():
+    ''' Check behavior when pandas input, ndarray desired output, column processing needed '''
+    i = LabelSwapInjector()
+    data = pd.DataFrame({'a':[0,0], 'b':[1,2]})
+    copy, cols = i._preprocess(data, 'a', 'b')
+    expected_copy = np.array([[0,1], [0,2]])
+    assert np.array_equal(copy, expected_copy)
+    assert cols == (0,1)
+
+def test_preprocess_4():
+    ''' Check behavior when pandas input, pandas desired output, column processing needed '''
+    i = LabelSwapInjector()
+    data = pd.DataFrame({'a':[0,0], 'b':[1,2]})
+    copy, cols = i._preprocess(data, 'a', 'b', return_df=True)
+    assert copy.equals(pd.DataFrame({'a':[0,0], 'b':[1,2]}))
+    assert cols == ('a', 'b')
+
+def test_preprocess_5():
+    ''' Check preprocessing failure on non-DF/ndarray data '''
+    i = LabelSwapInjector()
+    with pytest.raises(ValueError):
+        i._preprocess([[1,2], [2,3]], 'col_0')
+
+def test_call_1():
+    ''' Check Injector ABC has no call implementation '''
+    i = Injector()
+    with pytest.raises(NotImplementedError):
+        i()
+
+def test_postprocess_1():
+    ''' Check behavior when initial data pandas, input data numpy'''
+    i = LabelSwapInjector()
+    i._columns = ['a', 'b']
+    inp = np.ones((2,2))
+    exp = pd.DataFrame({'a':[1.0,1.0], 'b':[1.0,1.0]})
+    assert i._postprocess(inp).equals(exp)
+
+def test_postprocess_2():
+    ''' Check behavior when initial data ndarray, input data dataframe '''
+    i = LabelSwapInjector()
+    i._columns = None
+    inp = pd.DataFrame({'a':[1.0,1.0], 'b':[1.0,1.0]})
+    assert np.array_equal(i._postprocess(inp), np.ones((2,2)))
+
+def test_postprocess_3():
+    ''' Check behavior when initial data pandas, input data pandas '''
+    i = LabelSwapInjector()
+    i._columns = ['a', 'b']
+    inp = pd.DataFrame({'a':[1.0,1.0], 'b':[1.0,1.0]})
+    assert i._postprocess(inp).equals(inp)
+
+def test_postprocess_4():
+    ''' Check behavior when initial data numpy, input data numpy '''
+    i = LabelSwapInjector()
+    i._columns = None
+    inp = np.ones((2,2))
+    assert np.array_equal(inp, i._postprocess(inp))
\ No newline at end of file
diff --git a/tests/menelaus/injection/test_label_manipulation.py b/tests/menelaus/injection/test_label_manipulation.py
new file mode 100644
index 00000000..dba337e9
--- /dev/null
+++ b/tests/menelaus/injection/test_label_manipulation.py
@@ -0,0 +1,85 @@
+import pytest
+import numpy as np
+
+
+from menelaus.injection import LabelJoinInjector, LabelSwapInjector, LabelProbabilityInjector, LabelDirichletInjector
+
+
+def test_join_1():
+    ''' Check correct join behavior '''
+    i = LabelJoinInjector()
+    data = np.array([[0,5], [0,5], [1,5], [1,5]])
+    expected_data = np.array([[0,5], [2,5], [2,5], [1,5]])
+    copy = i(data, from_index=1, to_index=3, target_col=0, class_1=0, class_2=1, new_class=2)
+    assert np.array_equal(copy, expected_data)
+
+def test_swap_1():
+    ''' Check correct swap in pd.DataFrame '''
+    i = LabelSwapInjector()
+    data = np.array([[0,5], [0,5], [1,6], [1,6]])
+    expected_data = np.array([[1,5], [1,5], [0,6], [0,6]])
+    copy = i(data, from_index=0, to_index=4, target_col=0, class_1=0, class_2=1)
+    assert np.array_equal(copy, expected_data)
+
+def test_probability_shift_1():
+    ''' Ensure probability shift causes some drift in numpy.ndarray '''
+    i = LabelProbabilityInjector()
+    data = np.array([
+        [0,0,0], [1,0,0], [2,0,0], [3,0,0], [4,0,0], [5,0,0],
+        [6,0,0], [7,0,0], [8,0,0], [9,0,0]
+    ])
+    probs = {k:0.1 for k in range(0,10)}
+    np.random.seed(0)
+    new_data = i(data, from_index=1, to_index=9, target_col=0, class_probabilities=probs)
+    assert np.array_equal(data[0], new_data[0])
+    assert np.array_equal(data[-1], new_data[-1])
+    assert not np.array_equal(data[1:9], new_data[1:9])
+
+def test_probability_shift_2():
+    ''' Check drift injected when not all classes specified '''
+    i = LabelProbabilityInjector()
+    data = np.array([
+        [0,0,0], [1,0,0], [2,0,0], [3,0,0], [4,0,0], [5,0,0],
+        [6,0,0], [7,0,0], [8,0,0], [9,0,0]
+    ])
+    probs = {k:0.075 for k in range(2,10)}
+    np.random.seed(0)
+    _ = i(data, from_index=0, to_index=10, target_col=0, class_probabilities=probs)
+    # individual = class prob / class size, + leftover prob from unspecified classes
+    # 1-(8*0.075)=0.4, classes 0, 1 will have prob 0.2 each
+    # class sizes are 1, expected distr is simple
+    expected_distr = [0.2, 0.2, 0.075, 0.075, 0.075, 0.075, 0.075, 0.075, 0.075, 0.075]
+    p_distr = np.round(i._p_distribution, decimals=3)
+    assert np.array_equal(p_distr, expected_distr)
+
+def test_probability_shift_4():
+    ''' Check ValueError when probabilities add to >1 '''
+    i = LabelProbabilityInjector()
+    data = np.random.randint(0,5,(10,3))
+    with pytest.raises(ValueError):
+        class_1 = np.unique(data[:,0])[0]
+        class_2 = np.unique(data[:,0])[1]
+        probs = {class_1: 1.0, class_2: 0.5}
+        i(data, from_index=0, to_index=3, target_col=0, class_probabilities=probs)
+
+def test_probability_shift_5():
+    ''' Check ValueError when given probabilities for classes not in data '''
+    i = LabelProbabilityInjector()
+    data = np.random.randint(0,5,(10,3))
+    with pytest.raises(ValueError):
+        probs = {1000: 0.5}
+        i(data, from_index=0, to_index=3, target_col=0, class_probabilities=probs)
+
+def test_dirichlet_shift_1():
+    ''' Ensure Dirichlet shift causes some drift '''
+    i = LabelDirichletInjector()
+    data = np.array([
+        [0,0,0], [1,0,0], [2,0,0], [3,0,0], [4,0,0], [5,0,0],
+        [6,0,0], [7,0,0], [8,0,0], [9,0,0]
+    ])
+    alpha = {k:1 for k in range(0,10)}
+    np.random.seed(0)
+    new_data = i(data, from_index=1, to_index=9, target_col=0, alpha=alpha)
+    assert np.array_equal(data[0], new_data[0])
+    assert np.array_equal(data[-1], new_data[-1])
+    assert not np.array_equal(data[1:9], new_data[1:9])
\ No newline at end of file
diff --git a/tests/menelaus/injection/test_noise.py b/tests/menelaus/injection/test_noise.py
index 8bca2a44..462b1177 100644
--- a/tests/menelaus/injection/test_noise.py
+++ b/tests/menelaus/injection/test_noise.py
@@ -1,36 +1,41 @@
 """ Module to test noise injection functions """
 
-import pytest
 import pandas as pd
 import numpy as np
 
-from menelaus.injection import brownian_noise, random_walk
+from menelaus.injection import BrownianNoiseInjector
 
 
 def test_brownian_noise_1():
-    ''' Check general behavior for pandas DataFrame'''
-    old_values = [0,1,2,3,4,5,6]
-    data = pd.DataFrame({'a': old_values})
-    data = brownian_noise(data, 'a', 0, 1, 5)
-    assert not data.equals(pd.Series(old_values))
-
-def test_brownian_noise_2():
-    ''' Check general behavior for numpy array '''
-    old_values = np.array([[0],[1],[2],[3],[4],[5],[6]])
-    data = np.copy(old_values)
-    data = brownian_noise(data, 0, 0, 1, 5)
-    assert not np.array_equal(data, old_values)
-
-def test_bronwian_noise_3():
-    ''' Check ValueError if data neither pandas DataFrame nor numpy array '''
-    data = [[1,2], [0,1]]
-    with pytest.raises(ValueError):
-        brownian_noise(data, 0, 0, 0, 2)
+    ''' Check RW noise injected correctly '''
+    i = BrownianNoiseInjector()
+    exp_data = np.array([[0],[1],[2],[3],[4],[5],[6]])
+    copy = i(exp_data, from_index=1, to_index=5, col=0, x0=0, random_state=0)
+    
+    # copy of random walk function
+    np.random.seed(0)
+    w = np.ones(5-1) * 0
+    for i in range(1,5-1):
+        yi = np.random.choice([1,-1])
+        w[i] = w[i-1] + (yi / np.sqrt(5-1))
+    
+    exp_data[1:5, 0] = exp_data[1:5, 0] + w
+
+    # XXX - repeated code for testing random walk, no good alternatives
+    assert np.array_equal(exp_data, copy)
 
 def test_random_walk():
-    ''' Check random walk generally behaves within bounds '''
-    vec = random_walk(10, 0)
-    for i in range(1,10):
-        assert vec[i] <= vec[i-1] + 1/np.sqrt(10)
-        assert vec[i] >= vec[i-1] - 1/np.sqrt(10)
+    ''' Check random walk values '''
+    vec = BrownianNoiseInjector._random_walk(steps=5, x0=0, random_state=0)
+
+    # copy of random walk function
+    np.random.seed(0)
+    w = np.ones(5) * 0
+    for i in range(1, 5):
+        yi = np.random.choice([1, -1])
+        w[i] = w[i - 1] + (yi / np.sqrt(5))
+    
+    # XXX - repeating code for test coverage, although this may help if
+    #       implementation of random walk changes over time
+    assert np.array_equal(w, vec)