def test_len_t_should_return_the_number_of_records_where_treatment_equals_0(self):
df = pd.DataFrame(data=np.random.rand(6, 2), columns=["var1", "var2"])
df["Treatment"] = [random.sample(range(2), 1)[0] for i in range(6)]
length = df[df["Treatment"] == 0].shape[0]
result = utils.len_t(df, treatment=0)
self.assertEqual(length, result)
def test_len_t_should_return_the_number_of_records_where_treatment_equals_0_and_treatment_col_is_not_default(
self):
df = pd.DataFrame(data=np.random.rand(6, 2), columns=['var1', 'var2'])
df['MarketedTo'] = [random.sample(range(2), 1)[0] for i in range(6)]
length = df[df['MarketedTo'] == 0].shape[0]
result = utils.len_t(df, treatment=0, col_treatment='MarketedTo')
self.assertEqual(length, result)