From 235b49f3c3ed441776c10d38b8ecf74501e05e56 Mon Sep 17 00:00:00 2001
From: Sebastian Villena <97059974+ruisebas@users.noreply.github.com>
Date: Mon, 8 Apr 2024 12:08:26 -0400
Subject: [PATCH] fix(DataStore): Fixing a crash when attempting to create a
 model with a predicate. (#3600)

---
 .../Storage/StorageEngine.swift               |  1 +
 .../Storage/StorageEngineTestsHasOne.swift    | 22 +++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/StorageEngine.swift b/AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/StorageEngine.swift
index 078bf60624..53c213e4fb 100644
--- a/AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/StorageEngine.swift
+++ b/AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/StorageEngine.swift
@@ -206,6 +206,7 @@ final class StorageEngine: StorageEngineBehavior {
                 "Cannot apply a condition on model which does not exist.",
                 "Save the model instance without a condition first.")
             completion(.failure(causedBy: dataStoreError))
+            return
         }
 
         do {
diff --git a/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Storage/StorageEngineTestsHasOne.swift b/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Storage/StorageEngineTestsHasOne.swift
index 14133c3701..2627faa3a0 100644
--- a/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Storage/StorageEngineTestsHasOne.swift
+++ b/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Storage/StorageEngineTestsHasOne.swift
@@ -70,6 +70,28 @@ class StorageEngineTestsHasOne: StorageEngineTestsBase {
         }
     }
 
+    /// Given: A model that does not exist
+    /// When: save is called with a predicate
+    /// Then: A DataStoreError.invalidCondition error is returned
+    func testSaveModelWithPredicate_shouldFail() {
+        let team = Team(name: "Team")
+        let saveFinished = expectation(description: "Save finished")
+        storageEngine.save(team, condition: Team.keys.name.beginsWith("T")) { result in
+            defer {
+                saveFinished.fulfill()
+            }
+            guard case .failure(let error) = result,
+                  case . invalidCondition(let errorDescription, let recoverySuggestion, _) = error else {
+                XCTFail("Expected failure with .invalidCondition, got \(result)")
+                return
+            }
+
+            XCTAssertEqual(errorDescription, "Cannot apply a condition on model which does not exist.")
+            XCTAssertEqual(recoverySuggestion, "Save the model instance without a condition first.")
+        }
+        wait(for: [saveFinished], timeout: defaultTimeout)
+    }
+
     func testBelongsToRelationshipWithoutOwner() {
         let teamA = Team(name: "A-Team")
         let projectA = Project(name: "ProjectA", team: teamA)