From e95ce2d390ec7bb452e26c78a121c07a7bd31e90 Mon Sep 17 00:00:00 2001
From: Aaron Bach <bachya1208@gmail.com>
Date: Fri, 8 Mar 2024 20:27:56 -0700
Subject: [PATCH] Make sure Notion saves new refresh token upon startup
 (#112676)

* Make sure Notion saves new refresh token upon startup

* Code review

* Typing

* Smoother syntax

* Fix tests

* Fix tests for real
---
 homeassistant/components/notion/__init__.py |  4 ++++
 tests/components/notion/conftest.py         |  7 ++++---
 tests/components/notion/test_config_flow.py | 16 ++++++++++++----
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/notion/__init__.py b/homeassistant/components/notion/__init__.py
index 3ed2c7bdb930..33a749909282 100644
--- a/homeassistant/components/notion/__init__.py
+++ b/homeassistant/components/notion/__init__.py
@@ -180,6 +180,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     # Create a callback to save the refresh token when it changes:
     entry.async_on_unload(client.add_refresh_token_callback(async_save_refresh_token))
 
+    # Save the client's refresh token if it's different than what we already have:
+    if (token := client.refresh_token) and token != entry.data[CONF_REFRESH_TOKEN]:
+        async_save_refresh_token(token)
+
     hass.config_entries.async_update_entry(entry, **entry_updates)
 
     async def async_update() -> NotionData:
diff --git a/tests/components/notion/conftest.py b/tests/components/notion/conftest.py
index 61a54ca58bad..366f78b8c6c1 100644
--- a/tests/components/notion/conftest.py
+++ b/tests/components/notion/conftest.py
@@ -9,8 +9,8 @@ from aionotion.sensor.models import Sensor
 from aionotion.user.models import UserPreferences
 import pytest
 
-from homeassistant.components.notion import DOMAIN
-from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
+from homeassistant.components.notion import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN
+from homeassistant.const import CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
 from tests.common import MockConfigEntry, load_fixture
@@ -81,7 +81,8 @@ def config_fixture():
     """Define a config entry data fixture."""
     return {
         CONF_USERNAME: TEST_USERNAME,
-        CONF_PASSWORD: TEST_PASSWORD,
+        CONF_USER_UUID: TEST_USER_UUID,
+        CONF_REFRESH_TOKEN: TEST_REFRESH_TOKEN,
     }
 
 
diff --git a/tests/components/notion/test_config_flow.py b/tests/components/notion/test_config_flow.py
index 646bd7a6e87d..72bb3dfee0ba 100644
--- a/tests/components/notion/test_config_flow.py
+++ b/tests/components/notion/test_config_flow.py
@@ -10,7 +10,7 @@ from homeassistant.config_entries import SOURCE_REAUTH, SOURCE_USER
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from .conftest import TEST_REFRESH_TOKEN, TEST_USER_UUID, TEST_USERNAME
+from .conftest import TEST_PASSWORD, TEST_REFRESH_TOKEN, TEST_USER_UUID, TEST_USERNAME
 
 pytestmark = pytest.mark.usefixtures("mock_setup_entry")
 
@@ -26,7 +26,6 @@ pytestmark = pytest.mark.usefixtures("mock_setup_entry")
 async def test_create_entry(
     hass: HomeAssistant,
     client,
-    config,
     errors,
     get_client_with_exception,
     mock_aionotion,
@@ -44,13 +43,22 @@ async def test_create_entry(
         get_client_with_exception,
     ):
         result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_USER}, data=config
+            DOMAIN,
+            context={"source": SOURCE_USER},
+            data={
+                CONF_USERNAME: TEST_USERNAME,
+                CONF_PASSWORD: TEST_PASSWORD,
+            },
         )
         assert result["type"] == data_entry_flow.FlowResultType.FORM
         assert result["errors"] == errors
 
     result = await hass.config_entries.flow.async_configure(
-        result["flow_id"], user_input=config
+        result["flow_id"],
+        user_input={
+            CONF_USERNAME: TEST_USERNAME,
+            CONF_PASSWORD: TEST_PASSWORD,
+        },
     )
     assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
     assert result["title"] == TEST_USERNAME