import unittest
import pandas as pd
from nnfwtbn import Process, Cut
[docs]class ProcessTestCase(unittest.TestCase):
"""
Test the implementation of the Process class.
"""
[docs] def test_init(self):
"""
Check that all arguments are stored internally.
"""
process = Process("Top", selection=lambda d: d.is_top)
self.assertEqual(process.label, "Top")
self.assertIsNotNone(process.selection)
[docs] def test_no_type(self):
"""
Check that processes don't accept the 'type' argument.
"""
self.assertRaises(TypeError, Process, "Top", type="fill",
selection=lambda d: d.is_top)
[docs] def test_no_color(self):
"""
Check that processes don't accept the 'color' argument.
"""
self.assertRaises(TypeError, Process, "Top", color="#ff0000",
selection=lambda d: d.is_top)
[docs] def test_lambda(self):
"""
Check that a lambda selection is converted to a cut object.
"""
process = Process("Top", selection=lambda d: d.is_top)
self.assertEqual(process.label, "Top")
self.assertIsInstance(process.selection, Cut)
[docs] def test_init_selection_and_range(self):
"""
Check that an error is raised if both selection and range is passed to
the constructor.
"""
self.assertRaises(ValueError, Process, "Top",
selection=lambda d: d.is_top, range=(-10, -7))
[docs] def test_no_selection(self):
"""
Check that an issue is raised if no selection method is used.
"""
self.assertRaises(ValueError, Process, "Top")
[docs] def test_repr_cut(self):
"""
Check that the expected string representation is returned when the
object uses a selection cut.
"""
process = Process("Top", selection=lambda d: d.is_top)
self.assertEqual(repr(process), "<Process 'Top': (func)>")
[docs] def test_repr_range(self):
"""
Check that the expected string representation is returned when the
object uses a range selection.
"""
process = Process("Top", range=(5, 10))
self.assertEqual(repr(process), "<Process 'Top': [5, 10]>")
[docs] def test_default_range_var(self):
"""
Check that the default range_var value is taken from the class
property.
"""
old_default = Process.DEFAULT_RANGE_VAR
Process.DEFAULT_RANGE_VAR = 'some_thing_other'
process = Process("Top", range=(5, 10))
self.assertEqual(process.range_var, "some_thing_other")
Process.DEFAULT_RANGE_VAR = old_default
[docs] def test_range_var_store(self):
"""
Check that the range_var argument takes precedence over the class
constant.
"""
process = Process("Top", range=(5, 10), range_var="yet_another")
self.assertEqual(process.range_var, "yet_another")
[docs] def test_range_type(self):
"""
Check that an exception is raised if the range argument is not a tuple
of two numbers.
"""
self.assertRaises(ValueError, Process, "Top", range=[5, 10])
self.assertRaises(ValueError, Process, "Top", range=(5, 10, 20))
self.assertRaises(ValueError, Process, "Top", range=(5, ))
self.assertRaises(ValueError, Process, "Top", range=3)
[docs] def generate_df(self):
"""
Generate a toy dataframe.
"""
return pd.DataFrame({'process_id': [1, 2, 1],
'momentum_t': [100, 90, 110]})
[docs] def test_call(self):
"""
Check that calling a process returns a dataframe with selected events.
"""
process = Process("Signal", range=(1, 1), range_var="process_id")
df = self.generate_df()
process_df = process(df)
self.assertEqual(list(process_df.momentum_t), [100, 110])
[docs] def test_idx_array(self):
"""
Check that calling idx_array() returns an index array with selected events.
"""
process = Process("Signal", range=(1, 1), range_var="process_id")
df = self.generate_df()
selected = process.idx_array(df)
self.assertEqual(list(selected), [True, False, True])