From 823194a6dd176595508344a0985204f17dcd22c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20B=C4=85k?= <bakmatthew@icloud.com>
Date: Wed, 26 Jun 2024 10:27:53 +0200
Subject: [PATCH] Support variable declaration with forced unwrapped type

---
 Examples/Sources/ViewModel.swift                  |  1 +
 .../VariablesImplementationFactory.swift          | 10 +++++++++-
 .../Factories/UT_SpyFactory.swift                 | 15 +++++++++++++++
 .../UT_VariablesImplementationFactory.swift       |  7 +++++++
 4 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/Examples/Sources/ViewModel.swift b/Examples/Sources/ViewModel.swift
index c6b2141..0797ebf 100644
--- a/Examples/Sources/ViewModel.swift
+++ b/Examples/Sources/ViewModel.swift
@@ -5,6 +5,7 @@ protocol ServiceProtocol {
   var name: String { get }
   var anyProtocol: any Codable { get set }
   var secondName: String? { get }
+  var address: String! { get }
   var added: () -> Void { get set }
   var removed: (() -> Void)? { get set }
 
diff --git a/Sources/SpyableMacro/Factories/VariablesImplementationFactory.swift b/Sources/SpyableMacro/Factories/VariablesImplementationFactory.swift
index fdeff47..402ec73 100644
--- a/Sources/SpyableMacro/Factories/VariablesImplementationFactory.swift
+++ b/Sources/SpyableMacro/Factories/VariablesImplementationFactory.swift
@@ -53,8 +53,16 @@ struct VariablesImplementationFactory {
       // Since the count of `bindings` is exactly 1, it is safe to force unwrap it.
       let binding = protocolVariableDeclaration.bindings.first!
 
-      if binding.typeAnnotation?.type.is(OptionalTypeSyntax.self) == true {
+      /*
+       var name: String?
+       var name: String!
+       */
+      if binding.typeAnnotation?.type.is(OptionalTypeSyntax.self) == true ||
+          binding.typeAnnotation?.type.is(ImplicitlyUnwrappedOptionalTypeSyntax.self) == true {
         accessorRemovalVisitor.visit(protocolVariableDeclaration)
+      /*
+       var name: String
+       */
       } else {
         try protocolVariableDeclarationWithGetterAndSetter(binding: binding)
 
diff --git a/Tests/SpyableMacroTests/Factories/UT_SpyFactory.swift b/Tests/SpyableMacroTests/Factories/UT_SpyFactory.swift
index fe080e5..a63d730 100644
--- a/Tests/SpyableMacroTests/Factories/UT_SpyFactory.swift
+++ b/Tests/SpyableMacroTests/Factories/UT_SpyFactory.swift
@@ -413,6 +413,21 @@ final class UT_SpyFactory: XCTestCase {
     )
   }
 
+  func testDeclarationForcedUnwrappedVariable() throws {
+    try assertProtocol(
+      withDeclaration: """
+        protocol ServiceProtocol {
+            var data: String! { get set }
+        }
+        """,
+      expectingClassDeclaration: """
+        class ServiceProtocolSpy: ServiceProtocol {
+            var data: String!
+        }
+        """
+    )
+  }
+
   func testDeclarationExistentialVariable() throws {
     try assertProtocol(
       withDeclaration: """
diff --git a/Tests/SpyableMacroTests/Factories/UT_VariablesImplementationFactory.swift b/Tests/SpyableMacroTests/Factories/UT_VariablesImplementationFactory.swift
index ee49aa6..1a0304a 100644
--- a/Tests/SpyableMacroTests/Factories/UT_VariablesImplementationFactory.swift
+++ b/Tests/SpyableMacroTests/Factories/UT_VariablesImplementationFactory.swift
@@ -31,6 +31,13 @@ final class UT_VariablesImplementationFactory: XCTestCase {
     )
   }
 
+  func testVariablesDeclarationsForcedUnwrapped() throws {
+    try assertProtocolVariable(
+      withVariableDeclaration: "var foo: String! { get }",
+      expectingVariableDeclaration: "var foo: String!"
+    )
+  }
+
   func testVariablesDeclarationsClosure() throws {
     try assertProtocolVariable(
       withVariableDeclaration: "var completion: () -> Void { get }",