From 8b0bee5e8bc724e3942e71f073694d85acadc561 Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Sun, 19 Jun 2022 21:53:09 -0700
Subject: [PATCH] [fbprophet] Fix frequencies (#20326)

---
 superset/utils/pandas_postprocessing/utils.py |  8 +--
 .../pandas_postprocessing/test_prophet.py     | 61 +++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/superset/utils/pandas_postprocessing/utils.py b/superset/utils/pandas_postprocessing/utils.py
index 46f5a0c505..ab39135918 100644
--- a/superset/utils/pandas_postprocessing/utils.py
+++ b/superset/utils/pandas_postprocessing/utils.py
@@ -86,10 +86,10 @@ PROPHET_TIME_GRAIN_MAP = {
     "P1M": "M",
     "P3M": "Q",
     "P1Y": "A",
-    "1969-12-28T00:00:00Z/P1W": "W",
-    "1969-12-29T00:00:00Z/P1W": "W",
-    "P1W/1970-01-03T00:00:00Z": "W",
-    "P1W/1970-01-04T00:00:00Z": "W",
+    "1969-12-28T00:00:00Z/P1W": "W-SUN",
+    "1969-12-29T00:00:00Z/P1W": "W-MON",
+    "P1W/1970-01-03T00:00:00Z": "W-SAT",
+    "P1W/1970-01-04T00:00:00Z": "W-SUN",
 }
 
 RESAMPLE_METHOD = ("asfreq", "bfill", "ffill", "linear", "median", "mean", "sum")
diff --git a/tests/unit_tests/pandas_postprocessing/test_prophet.py b/tests/unit_tests/pandas_postprocessing/test_prophet.py
index e4f3ed8cfc..6da3a7a591 100644
--- a/tests/unit_tests/pandas_postprocessing/test_prophet.py
+++ b/tests/unit_tests/pandas_postprocessing/test_prophet.py
@@ -17,6 +17,7 @@
 from datetime import datetime
 from importlib.util import find_spec
 
+import pandas as pd
 import pytest
 
 from superset.exceptions import InvalidPostProcessingError
@@ -50,6 +51,66 @@ def test_prophet_valid():
     assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 5, 31)
     assert len(df) == 9
 
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 2), datetime(2022, 1, 9)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="P1W",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 16)
+    assert len(df) == 3
+
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 2), datetime(2022, 1, 9)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="1969-12-28T00:00:00Z/P1W",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 16)
+    assert len(df) == 3
+
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 3), datetime(2022, 1, 10)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="1969-12-29T00:00:00Z/P1W",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 17)
+    assert len(df) == 3
+
+    df = prophet(
+        df=pd.DataFrame(
+            {
+                "__timestamp": [datetime(2022, 1, 8), datetime(2022, 1, 15)],
+                "x": [1, 1],
+            }
+        ),
+        time_grain="P1W/1970-01-03T00:00:00Z",
+        periods=1,
+        confidence_interval=0.9,
+    )
+
+    assert df[DTTM_ALIAS].iloc[-1].to_pydatetime() == datetime(2022, 1, 22)
+    assert len(df) == 3
+
 
 def test_prophet_valid_zero_periods():
     pytest.importorskip("prophet")