Kpi forecasting: fix tests (#222)

* replaced black with ruff and updated files

* ruff updates

* update circle ci configs

* reverted configs

.circleci/config.yml

@@ -234,7 +234,7 @@ jobs:
           command: docker build -t app:build jobs/kpi-forecasting/
       - run:
           name: Test Code
-          command: docker run app:build black .
+          command: docker run app:build pytest --ruff --ruff-format
 
 
   build-job-mozaggregator2bq:

jobs/kpi-forecasting/ci_job.yaml

@@ -12,4 +12,4 @@ build-job-kpi-forecasting:
         command: docker build -t app:build jobs/kpi-forecasting/
     - run:
         name: Test Code
-        command: docker run app:build black .
+        command: docker run app:build pytest --ruff --ruff-format

jobs/kpi-forecasting/kpi_forecasting/configs/model_inputs/__init__.py

@@ -1,8 +1,7 @@
 import attr
-from typing import List, Dict, Optional, Union
+from typing import List, Optional, Union
 from pathlib import Path
 
-import pandas as pd
 
 from kpi_forecasting.inputs import YAML
 

jobs/kpi-forecasting/kpi_forecasting/metric_hub.py

@@ -1,12 +1,10 @@
 import pandas as pd
 
 from dataclasses import dataclass
-from datetime import datetime
 from dotmap import DotMap
 from google.cloud import bigquery
 from mozanalysis.config import ConfigLoader
 from textwrap import dedent
-from typing import Dict
 
 from kpi_forecasting.utils import parse_end_date
 

jobs/kpi-forecasting/kpi_forecasting/models/base_forecast.py

@@ -4,12 +4,12 @@ import pandas as pd
 
 from google.cloud import bigquery
 from google.cloud.bigquery.enums import SqlTypeNames as bq_types
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 from kpi_forecasting import pandas_extras as pdx
 from kpi_forecasting.metric_hub import MetricHub
 from pandas.api import types as pd_types
-from typing import Dict, List, Tuple
+from typing import Dict, List
 
 
 @dataclass
@@ -166,7 +166,8 @@ class BaseForecast:
             # forecast is generated in the middle of the month.
             .add(overlap[["value"]].values)
             # calculate summary values, aggregating by submission_date,
-            .agg(aggregations, axis=1).reset_index()
+            .agg(aggregations, axis=1)
+            .reset_index()
             # "melt" the df from wide-format to long-format.
             .melt(id_vars="submission_date", var_name="measure")
         )

jobs/kpi-forecasting/kpi_forecasting/models/funnel_forecast.py

@@ -222,7 +222,6 @@ class FunnelForecast(BaseForecast):
 
         # build training dataframe
         if task == "train":
-
             # find indices in observed_df for rows that exactly match segment dict
             segment_historical_indices = (
                 self.observed_df[list(segment_settings.segment)]
@@ -532,7 +531,8 @@ class FunnelForecast(BaseForecast):
             # forecast is generated in the middle of the month.
             .add(overlap[["value"]].values)
             # calculate summary values, aggregating by submission_date,
-            .agg(aggregations, axis=1).reset_index()
+            .agg(aggregations, axis=1)
+            .reset_index()
         ).rename(columns=self._percentile_name_map(percentiles))
 
         # add datasource-specific metadata columns

jobs/kpi-forecasting/kpi_forecasting/tests/test_funnel_forecast.py

@@ -1,2 +0,0 @@
-from kpi_forecasting.models import funnel_forecast
-from kpi_forecasting.metric_hub import MetricHub

jobs/kpi-forecasting/kpi_forecasting/tests/test_metric_hub.py

@@ -1,4 +1,5 @@
-from datetime import datetime
+from datetime import datetime, timezone
+import re
 
 from pandas import to_datetime
 from kpi_forecasting.metric_hub import MetricHub
@@ -11,12 +12,12 @@ def test_metrichub_for_dau_kpi():
         slug="mobile_daily_active_users_v1",
         start_date="2024-01-01",
     )
-    now = to_datetime(datetime.utcnow()).date()
+    now = to_datetime(datetime.now(timezone.utc)).date()
 
     query = test_metric_hub.query()
-    query_where = f"WHERE submission_date BETWEEN '2024-01-01' AND '{now}'\n GROUP BY"
+    query_where = f"WHERE submission_date BETWEEN '2024-01-01' AND '{now}'\nGROUP BY"
 
-    assert query_where in query
+    assert re.sub(r"[\n\t\s]*", "", query_where) in re.sub(r"[\n\t\s]*", "", query)
     assert "\n    AND" not in query
 
 
@@ -29,7 +30,7 @@ def test_metrichub_with_where():
     )
 
     query = test_metric_hub.query()
-    assert f"\n    AND {test_metric_hub.where}" in query
+    assert f"\n    {test_metric_hub.where}" in query
 
 
 def test_metrichub_with_segments():
@@ -41,7 +42,10 @@ def test_metrichub_with_segments():
     )
 
     query = test_metric_hub.query()
-    assert "segment1 AS test_segment1,\n     segment2 AS test_segment2" in query
+    include_segment_no_whitespace = re.sub(
+        r"[\n\t\s]*", "", "segment1 AS test_segment1, segment2 AS test_segment2"
+    )
+    assert include_segment_no_whitespace in re.sub(r"[\n\t\s]*", "", query)
 
 
 def test_metrichub_with_segments_and_where():
@@ -54,8 +58,16 @@ def test_metrichub_with_segments_and_where():
     )
 
     query = test_metric_hub.query()
-    assert f"\n    AND {test_metric_hub.where}" in query
-    assert "segment1 AS test_segment1,\n     segment2 AS test_segment2" in query
+    query_no_whitespace = re.sub(r"[\n\t\s]*", "", query)
+    assert re.sub(r"[\n\t\s]*", "", test_metric_hub.where) in query_no_whitespace
+    assert (
+        re.sub(
+            r"[\n\t\s]*",
+            "",
+            "segment1 AS test_segment1,\n     segment2 AS test_segment2",
+        )
+        in query_no_whitespace
+    )
 
 
 def test_metrichub_no_end_date():
@@ -64,7 +76,7 @@ def test_metrichub_no_end_date():
         slug="mobile_daily_active_users_v1",
         start_date="2024-01-01",
     )
-    now = to_datetime(datetime.utcnow()).date()
+    now = to_datetime(datetime.now(timezone.utc)).date()
 
     assert test_metric_hub.end_date == now
 
@@ -76,7 +88,7 @@ def test_metrichub_last_complete_month():
         start_date="2024-01-01",
         end_date="last complete month",
     )
-    now = to_datetime(datetime.utcnow())
+    now = to_datetime(datetime.now(timezone.utc)).date()
     prev_date = previous_period_last_date("last complete month", now)
 
     assert test_metric_hub.end_date == to_datetime(prev_date).date()

jobs/kpi-forecasting/kpi_forecasting/tests/test_utils.py

@@ -4,7 +4,6 @@ from kpi_forecasting.utils import parse_end_date, previous_period_last_date
 
 
 def test_parse_end_date_with_date():
-
     dt = "2024-01-01"
     parsed_date = parse_end_date(dt)
 
@@ -12,7 +11,6 @@ def test_parse_end_date_with_date():
 
 
 def test_parse_end_date_with_none():
-
     dt = None
     parsed_date = parse_end_date(dt)
 
@@ -20,7 +18,6 @@ def test_parse_end_date_with_none():
 
 
 def test_parse_end_date_prev_year():
-
     dt = "last complete year"
     parsed_date = parse_end_date(dt)
 
@@ -30,7 +27,6 @@ def test_parse_end_date_prev_year():
 
 
 def test_parse_end_date_prev_month():
-
     dt = "last complete month"
     now = datetime.datetime(2024, 1, 1)
 

jobs/kpi-forecasting/kpi_forecasting/utils.py

@@ -2,7 +2,6 @@ import datetime
 
 
 def parse_end_date(end_date: str = None):
-
     dt = datetime.datetime.utcnow()
     if not end_date:
         return dt
@@ -12,7 +11,6 @@ def parse_end_date(end_date: str = None):
 
 
 def previous_period_last_date(last_period: str, now: datetime.datetime):
-
     if last_period not in ["last complete month", "last complete year"]:
         raise ValueError("Unrecognized end date string.")
 

jobs/kpi-forecasting/requirements.txt

@@ -56,7 +56,7 @@ pyasn1-modules==0.3.0
 PyMeeus==0.5.12
 pyparsing==3.0.9
 pytest==7.3.2
-pytest-black==0.3.12
+pytest-ruff==0.3.2
 python-dateutil==2.8.2
 pytz==2023.3
 PyYAML==6.0