def test_process_metrics(spark):
exp = Experiment('a-stub', '20190101', num_dates_enrollment=8)
enrollments = exp.get_enrollments(spark,
_get_enrollment_view(slug="a-stub"))
ds_df_A = register_data_source_fixture(spark, name='ds_df_A')
ds_df_B = register_data_source_fixture(spark, name='ds_df_B')
ds_A = DataSource.from_dataframe('ds_df_A', ds_df_A)
ds_B = DataSource.from_dataframe('ds_df_B', ds_df_B)
m1 = Metric.from_col('m1', ds_df_A.numeric_col, ds_A)
m2 = Metric.from_col('m2', ds_df_A.bool_col, ds_A)
m3 = Metric.from_col('m3', ds_df_B.numeric_col, ds_B)
metric_list = [m1, m2, m3]
exp = Experiment('a-stub', '20190101')
data_sources_and_metrics = exp._process_metrics(enrollments, metric_list)
assert len(data_sources_and_metrics) == 2
assert len(data_sources_and_metrics[ds_df_A]) == 2
assert len(data_sources_and_metrics[ds_df_B]) == 1
assert 'numeric_col' in repr(data_sources_and_metrics[ds_df_B][0])
assert '`m3`' in repr(data_sources_and_metrics[ds_df_B][0])
assert repr(data_sources_and_metrics[ds_df_B][0]) in {
"Column<b'numeric_col AS `m3`'>", # py3
"Column<numeric_col AS `m3`>", # py2
}
def test_process_metrics_dupe_data_source(spark):
exp = Experiment('a-stub', '20190101', num_dates_enrollment=8)
enrollments = exp.get_enrollments(spark,
_get_enrollment_view(slug="a-stub"))
ds_df = register_data_source_fixture(spark, name='ds_df_A')
ds_1 = DataSource.from_dataframe('ds_df_A', ds_df)
ds_2 = DataSource.from_dataframe('ds_df_A', ds_df)
m1 = Metric.from_col('m1', ds_df.numeric_col, ds_1)
m2 = Metric.from_col('m2', ds_df.bool_col, ds_2)
metric_list = [m1, m2]
exp = Experiment('a-stub', '20190101')
data_sources_and_metrics = exp._process_metrics(enrollments, metric_list)
assert len(data_sources_and_metrics) == 1
assert len(data_sources_and_metrics[ds_df]) == 2