From debcaf6fb7602db12f98cb5c02fbed4b5e7591dc Mon Sep 17 00:00:00 2001
From: Otto Winter <otto@otto-winter.com>
Date: Sat, 19 Feb 2022 15:47:50 +0100
Subject: [PATCH] Add ESP32C3 and ESP32S2 support to dashboard (#3152)

* Add ESP32C3 and ESP32S2 support to dashboard

* Format

* Fix tests
---
 esphome/components/esp32/__init__.py |  4 +--
 esphome/storage_json.py              | 12 ++++---
 esphome/wizard.py                    | 49 +++++++++++++++++++++-------
 tests/unit_tests/test_wizard.py      |  4 +--
 4 files changed, 50 insertions(+), 19 deletions(-)

diff --git a/esphome/components/esp32/__init__.py b/esphome/components/esp32/__init__.py
index 0b2c291ba2..478c9701b8 100644
--- a/esphome/components/esp32/__init__.py
+++ b/esphome/components/esp32/__init__.py
@@ -61,8 +61,8 @@ def set_core_data(config):
     return config
 
 
-def get_esp32_variant():
-    return CORE.data[KEY_ESP32][KEY_VARIANT]
+def get_esp32_variant(core_obj=None):
+    return (core_obj or CORE).data[KEY_ESP32][KEY_VARIANT]
 
 
 def only_on_variant(*, supported=None, unsupported=None):
diff --git a/esphome/storage_json.py b/esphome/storage_json.py
index 207a3edf57..d561de4ce6 100644
--- a/esphome/storage_json.py
+++ b/esphome/storage_json.py
@@ -64,7 +64,7 @@ class StorageJSON:
         # Web server port of the ESP, for example 80
         assert web_port is None or isinstance(web_port, int)
         self.web_port = web_port  # type: int
-        # The type of ESP in use, either ESP32 or ESP8266
+        # The type of hardware in use, like "ESP32", "ESP32C3", "ESP8266", etc.
         self.target_platform = target_platform  # type: str
         # The absolute path to the platformio project
         self.build_path = build_path  # type: str
@@ -99,6 +99,11 @@ class StorageJSON:
     def from_esphome_core(
         esph, old
     ):  # type: (CoreType, Optional[StorageJSON]) -> StorageJSON
+        hardware = esph.target_platform
+        if esph.is_esp32:
+            from esphome.components import esp32
+
+            hardware = esp32.get_esp32_variant(esph)
         return StorageJSON(
             storage_version=1,
             name=esph.name,
@@ -107,15 +112,14 @@ class StorageJSON:
             src_version=1,
             address=esph.address,
             web_port=esph.web_port,
-            target_platform=esph.target_platform,
+            target_platform=hardware,
             build_path=esph.build_path,
             firmware_bin_path=esph.firmware_bin,
             loaded_integrations=list(esph.loaded_integrations),
         )
 
     @staticmethod
-    def from_wizard(name, address, esp_platform):
-        # type: (str, str, str) -> StorageJSON
+    def from_wizard(name: str, address: str, esp_platform: str) -> "StorageJSON":
         return StorageJSON(
             storage_version=1,
             name=name,
diff --git a/esphome/wizard.py b/esphome/wizard.py
index c64ad3a583..34930ff66f 100644
--- a/esphome/wizard.py
+++ b/esphome/wizard.py
@@ -67,6 +67,27 @@ esp32:
     type: arduino
 """
 
+ESP32S2_CONFIG = """
+esp32:
+  board: {board}
+  framework:
+    type: esp-idf
+"""
+
+ESP32C3_CONFIG = """
+esp32:
+  board: {board}
+  framework:
+    type: esp-idf
+"""
+
+HARDWARE_BASE_CONFIGS = {
+    "ESP8266": ESP8266_CONFIG,
+    "ESP32": ESP32_CONFIG,
+    "ESP32S2": ESP32S2_CONFIG,
+    "ESP32C3": ESP32C3_CONFIG,
+}
+
 
 def sanitize_double_quotes(value):
     return value.replace("\\", "\\\\").replace('"', '\\"')
@@ -83,11 +104,7 @@ def wizard_file(**kwargs):
 
     config = BASE_CONFIG.format(**kwargs)
 
-    config += (
-        ESP8266_CONFIG.format(**kwargs)
-        if kwargs["platform"] == "ESP8266"
-        else ESP32_CONFIG.format(**kwargs)
-    )
+    config += HARDWARE_BASE_CONFIGS[kwargs["platform"]].format(**kwargs)
 
     config += LOGGER_API_CONFIG
 
@@ -119,16 +136,26 @@ def wizard_file(**kwargs):
 """
 
     # pylint: disable=consider-using-f-string
-    config += """
+    if kwargs["platform"] in ["ESP8266", "ESP32"]:
+        config += """
   # Enable fallback hotspot (captive portal) in case wifi connection fails
   ap:
     ssid: "{fallback_name}"
     password: "{fallback_psk}"
 
 captive_portal:
-""".format(
-        **kwargs
-    )
+    """.format(
+            **kwargs
+        )
+    else:
+        config += """
+  # Enable fallback hotspot in case wifi connection fails
+  ap:
+    ssid: "{fallback_name}"
+    password: "{fallback_psk}"
+    """.format(
+            **kwargs
+        )
 
     return config
 
@@ -147,10 +174,10 @@ def wizard_write(path, **kwargs):
         kwargs["platform"] = (
             "ESP8266" if board in esp8266_boards.ESP8266_BOARD_PINS else "ESP32"
         )
-    platform = kwargs["platform"]
+    hardware = kwargs["platform"]
 
     write_file(path, wizard_file(**kwargs))
-    storage = StorageJSON.from_wizard(name, f"{name}.local", platform)
+    storage = StorageJSON.from_wizard(name, f"{name}.local", hardware)
     storage_path = ext_storage_path(os.path.dirname(path), os.path.basename(path))
     storage.save(storage_path)
 
diff --git a/tests/unit_tests/test_wizard.py b/tests/unit_tests/test_wizard.py
index 59fcfbff60..79a5894075 100644
--- a/tests/unit_tests/test_wizard.py
+++ b/tests/unit_tests/test_wizard.py
@@ -3,14 +3,14 @@
 import esphome.wizard as wz
 import pytest
 from esphome.components.esp8266.boards import ESP8266_BOARD_PINS
-from mock import MagicMock
+from unittest.mock import MagicMock
 
 
 @pytest.fixture
 def default_config():
     return {
         "name": "test-name",
-        "platform": "test_platform",
+        "platform": "ESP8266",
         "board": "esp01_1m",
         "ssid": "test_ssid",
         "psk": "test_psk",