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