import pytest import os import pytest import numpy as np import pandas as pd from absl import logging from datetime import datetime from forecasting_toolkit.datastore.adapters.snowflake import ( SnowflakeDatasetAdapter ) from forecasting_toolkit.datastore.connectors.snowflake import ( snowflake_connector_factory, set_snowflake_environment ) from forecasting_toolkit.models.pipelines.inference.dataset.column_handlers import * @pytest.fixture def test_dataset_df(): dataset_table = "DATASET_STREAMS_DAILY_2022" datafile = "test_df.csv" limit = 10 filters = {} if os.path.exists(datafile): return pd.read_csv(datafile) with snowflake_connector_factory() as conn: # set snowflake environment logging.debug("setting up DB Env") # TODO: at a later point - makes this ENV vars that get injected set_snowflake_environment(conn_cursor=conn, warehouse="DEV_OWS_WAREHOUSE", database="DEV_ENGINEERING", schema = "AADAMU_DEBUT_FORECASTING_DBT") # create snowflake dataset adapter snowflake_dataset_adapter = SnowflakeDatasetAdapter(conn=conn) # fetch dataset logging.debug("Fetching dataset from snowflake") dataset_df = snowflake_dataset_adapter.fetch_dataset(snowflake_table=dataset_table, filters=filters, limit=limit) # persist dataset_df.to_csv(datafile, index=False) return dataset_df def test_rollfoward_snapshot_dates(test_dataset_df): """ tests to ensure snapshot date gets rolle forward""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) lst_dates = pd.to_datetime([datetime.now() + timedelta(days=x) for x in range(rollforward_days)]) new_df['SNAPSHOT_DATE'] = pd.to_datetime(new_df['SNAPSHOT_DATE']) compare_gen = zip(lst_dates,new_df['SNAPSHOT_DATE'].values) # assert to make sure its the same days assert np.allclose(a=[(t_1 - t_2).days for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_snapshot_year(test_dataset_df): """ tests to ensure snapshot year gets rolle forward""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) new_df = rollforward_snapshot_year(new_df) lst_dates = [datetime.now().year for _ in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['SNAPSHOT_YEAR'].values )) # check types assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) # assert to make sure it matches assert np.allclose(a=[(t_1 - t_2) for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_snapshot_isoweek(test_dataset_df): """ tests to ensure snapshot isoweek gets rolled forward""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) new_df = rollforward_snapshot_isoweek(new_df) lst_dates = [(datetime.now() + timedelta(days=x)).isocalendar()[1] for x in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['SNAPSHOT_ISOWEEK'].values )) # check types assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) # assert to make sure it matches assert np.allclose(a=[t_1 - t_2 for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_snapshot_month(test_dataset_df): """ tests to ensure snapshot month gets rolled forward""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) new_df = rollforward_snapshot_month(new_df) lst_dates = [(datetime.now() + timedelta(days=x)).date().month for x in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['SNAPSHOT_MONTH'].values )) # check types assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) # assert to make sure it matches assert np.allclose(a=[t_1 - t_2 for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_snapshot_day_of_year(test_dataset_df): """ tests to ensure snapshot day of year gets rolled forward""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) new_df = rollforward_snapshot_day_of_year(new_df) lst_dates = [(datetime.now() + timedelta(days=x)).timetuple().tm_yday for x in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['SNAPSHOT_DAY_OF_YEAR'].values)) # check types assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) # assert to make sure it matches assert np.allclose(a=[t_1 - t_2 for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_snapshot_day_of_week(test_dataset_df): """ tests to ensure snapshot day of weeks gets rolled forward""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) new_df = rollforward_snapshot_day_of_week(new_df) lst_dates = [(datetime.now() + timedelta(days=x)).weekday() for x in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['SNAPSHOT_DAY_OF_WEEK'].values)) # check types assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) # assert to make sure it matches assert np.allclose(a=[t_1 - t_2 for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_release_date(test_dataset_df): """ tests to ensure release date gets rolle forward""" rollforward_days = 10 new_df = rollforward_release_date(test_dataset_df, release_date=datetime.now()) lst_dates = pd.to_datetime([datetime.now() for _ in range(rollforward_days)]) compare_gen = iter(zip(lst_dates, new_df['RELEASE_DATE'].values)) # assert they are similar assert np.allclose(a=[(t_1 - t_2).days for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_release_year(test_dataset_df): """ tests to ensure release year is rolled forward properly""" rollforward_days = 10 new_df = rollforward_release_date(test_dataset_df, release_date=datetime.now()) new_df = rollforward_release_year(new_df) lst_dates = [datetime.now().year for _ in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['RELEASE_YEAR'].values)) # assert they are similar assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) assert np.allclose(a=[(t_1 - t_2) for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_release_month(test_dataset_df): """ tests to ensure release month is rolled forward properly""" rollforward_days = 10 new_df = rollforward_release_date(test_dataset_df, release_date=datetime.now()) new_df = rollforward_release_month(new_df) lst_dates = [datetime.now().month for _ in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['RELEASE_MONTH'].values)) # assert they are similar assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) assert np.allclose(a=[(t_1 - t_2) for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_release_day_of_week(test_dataset_df): """ tests to ensure release day of week is rolled forward properly""" rollforward_days = 10 new_df = rollforward_release_date(test_dataset_df, release_date=datetime.now()) new_df = rollforward_release_day_of_week(new_df) lst_dates = [datetime.now().weekday() for _ in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['RELEASE_DOW'].values)) # assert they are similar assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) assert np.allclose(a=[(t_1 - t_2) for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_release_weekiso(test_dataset_df): """ tests to ensure release isoweek is rolled forward properly""" rollforward_days = 10 new_df = rollforward_release_date(test_dataset_df, release_date=datetime.now()) new_df = rollforward_release_weekiso(new_df) lst_dates = [datetime.now().isocalendar()[1] for _ in range(rollforward_days)] compare_gen = iter(zip(lst_dates, new_df['RELEASE_WEEKISO'].values)) # assert they are similar assert isinstance(lst_dates, list) assert isinstance(lst_dates[0], int) assert np.allclose(a=[(t_1 - t_2) for (t_1, t_2) in compare_gen], b=np.zeros(len(lst_dates)), atol=0) def test_rollforward_release_age_days(test_dataset_df): """ tests to ensure release age in days""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) new_df = rollforward_release_date(new_df, release_date=datetime.now()) new_df = rollforward_release_age_days(new_df) compare_gen = iter(zip(new_df['SNAPSHOT_DATE'], new_df['RELEASE_DATE'].values)) release_age_days_lst = [(t_1 - t_2).days for (t_1,t_2) in compare_gen] print(new_df['SNAPSHOT_DATE'].values) # assert its a list assert isinstance(release_age_days_lst, list) assert isinstance(release_age_days_lst[0], int) # ensure we get snapshots from a day after release assert np.allclose(a=release_age_days_lst, b=new_df['RELEASE_AGE_DAYS'].values, atol=0) def test_rollforward_release_age_weeks(test_dataset_df): """ test rolling forward release age in weeks""" rollforward_days = 10 new_df = rollfoward_snapshot_dates(test_dataset_df, from_date= datetime.now(), rollforward_days=rollforward_days) new_df = rollforward_release_date(new_df, release_date=datetime.now()) new_df = rollforward_release_age_weeks(new_df) compare_gen = iter(zip(pd.to_datetime(new_df['SNAPSHOT_DATE']), pd.to_datetime(new_df['RELEASE_DATE'].values))) release_age_days_lst = [((t_1 - t_2).days)//7 for (t_1, t_2) in compare_gen] # print(new_df['SNAPSHOT_DATE'].values) # assert its a list assert isinstance(release_age_days_lst, list) assert isinstance(release_age_days_lst[0], int) # ensure we get snapshots from a day after release assert np.allclose(a=release_age_days_lst, b=new_df['RELEASE_AGE_WEEKS'].values)