From fbc84861c7fba05e40603b513c65d73a37f27136 Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Wed, 8 Dec 2021 09:22:03 +1300
Subject: [PATCH] Use new platform component config blocks for wizard (#2885)

---
 esphome/wizard.py               | 24 ++++++++++++++++++++++--
 tests/unit_tests/test_wizard.py |  9 +++++----
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/esphome/wizard.py b/esphome/wizard.py
index 5f4f347ba7..f2632caf71 100644
--- a/esphome/wizard.py
+++ b/esphome/wizard.py
@@ -45,9 +45,9 @@ OTA_BIG = r"""       ____ _______
 
 BASE_CONFIG = """esphome:
   name: {name}
-  platform: {platform}
-  board: {board}
+"""
 
+LOGGER_API_CONFIG = """
 # Enable logging
 logger:
 
@@ -55,6 +55,18 @@ logger:
 api:
 """
 
+ESP8266_CONFIG = """
+esp8266:
+  board: {board}
+"""
+
+ESP32_CONFIG = """
+esp32:
+  board: {board}
+  framework:
+    type: arduino
+"""
+
 
 def sanitize_double_quotes(value):
     return value.replace("\\", "\\\\").replace('"', '\\"')
@@ -71,6 +83,14 @@ def wizard_file(**kwargs):
 
     config = BASE_CONFIG.format(**kwargs)
 
+    config += (
+        ESP8266_CONFIG.format(**kwargs)
+        if kwargs["platform"] == "ESP8266"
+        else ESP32_CONFIG.format(**kwargs)
+    )
+
+    config += LOGGER_API_CONFIG
+
     # Configure API
     if "password" in kwargs:
         config += f"  password: \"{kwargs['password']}\"\n"
diff --git a/tests/unit_tests/test_wizard.py b/tests/unit_tests/test_wizard.py
index 18e040b0a6..59fcfbff60 100644
--- a/tests/unit_tests/test_wizard.py
+++ b/tests/unit_tests/test_wizard.py
@@ -11,7 +11,7 @@ def default_config():
     return {
         "name": "test-name",
         "platform": "test_platform",
-        "board": "test_board",
+        "board": "esp01_1m",
         "ssid": "test_ssid",
         "psk": "test_psk",
         "password": "",
@@ -105,6 +105,7 @@ def test_wizard_write_sets_platform(default_config, tmp_path, monkeypatch):
     If the platform is not explicitly set, use "ESP8266" if the board is one of the ESP8266 boards
     """
     # Given
+    del default_config["platform"]
     monkeypatch.setattr(wz, "write_file", MagicMock())
 
     # When
@@ -112,7 +113,7 @@ def test_wizard_write_sets_platform(default_config, tmp_path, monkeypatch):
 
     # Then
     generated_config = wz.write_file.call_args.args[1]
-    assert f"platform: {default_config['platform']}" in generated_config
+    assert "esp8266:" in generated_config
 
 
 def test_wizard_write_defaults_platform_from_board_esp8266(
@@ -132,7 +133,7 @@ def test_wizard_write_defaults_platform_from_board_esp8266(
 
     # Then
     generated_config = wz.write_file.call_args.args[1]
-    assert "platform: ESP8266" in generated_config
+    assert "esp8266:" in generated_config
 
 
 def test_wizard_write_defaults_platform_from_board_esp32(
@@ -152,7 +153,7 @@ def test_wizard_write_defaults_platform_from_board_esp32(
 
     # Then
     generated_config = wz.write_file.call_args.args[1]
-    assert "platform: ESP32" in generated_config
+    assert "esp32:" in generated_config
 
 
 def test_safe_print_step_prints_step_number_and_description(monkeypatch):