From 7046559c29cafec089a0c89636ea5d0388ff6c42 Mon Sep 17 00:00:00 2001
From: Jack Walker <jack@foresightmobile.com>
Date: Thu, 20 Jun 2024 11:11:56 +0100
Subject: [PATCH] Removed getIdentityFlags method and added traits variable to
 getFeatureFlags method

---
 FlagsmithClient/Classes/Flagsmith.swift      | 67 ++++++++++----------
 FlagsmithClient/Classes/FlagsmithError.swift |  4 ++
 2 files changed, 38 insertions(+), 33 deletions(-)

diff --git a/FlagsmithClient/Classes/Flagsmith.swift b/FlagsmithClient/Classes/Flagsmith.swift
index f61aaa4..c7a7054 100644
--- a/FlagsmithClient/Classes/Flagsmith.swift
+++ b/FlagsmithClient/Classes/Flagsmith.swift
@@ -83,36 +83,52 @@ public final class Flagsmith: @unchecked Sendable {
     ///   - identity: ID of the user (optional)
     ///   - completion: Closure with Result which contains array of Flag objects in case of success or Error in case of failure
     public func getFeatureFlags(forIdentity identity: String? = nil,
+                                traits: [Trait]? = nil,
                                 completion: @Sendable @escaping (Result<[Flag], any Error>) -> Void)
     {
         if let identity = identity {
-            getIdentity(identity) { result in
-                switch result {
-                case let .success(thisIdentity):
-                    completion(.success(thisIdentity.flags))
-                case let .failure(error):
-                    if self.defaultFlags.isEmpty {
-                        completion(.failure(error))
-                    } else {
-                        completion(.success(self.defaultFlags))
+            if let traits = traits {
+                apiManager.request(.postTraits(identity: identity, traits: traits)) { (result: Result<Traits, Error>) in
+                    switch result {
+                    case let .success(result):
+                        completion(.success(result.flags))
+                    case let .failure(error):
+                        self.handleFlagsError(error, completion: completion)
+                    }
+                }
+            } else {
+                getIdentity(identity) { result in
+                    switch result {
+                    case let .success(thisIdentity):
+                        completion(.success(thisIdentity.flags))
+                    case let .failure(error):
+                        self.handleFlagsError(error, completion: completion)
                     }
                 }
             }
         } else {
-            apiManager.request(.getFlags) { (result: Result<[Flag], Error>) in
-                switch result {
-                case let .success(flags):
-                    completion(.success(flags))
-                case let .failure(error):
-                    if self.defaultFlags.isEmpty {
-                        completion(.failure(error))
-                    } else {
-                        completion(.success(self.defaultFlags))
+            if let _ = traits {
+                completion(.failure(FlagsmithError.invalidArgument("You must provide an identity to set traits")))
+            } else {
+                apiManager.request(.getFlags) { (result: Result<[Flag], Error>) in
+                    switch result {
+                    case let .success(flags):
+                        completion(.success(flags))
+                    case let .failure(error):
+                        self.handleFlagsError(error, completion: completion)
                     }
                 }
             }
         }
     }
+    
+    private func handleFlagsError(_ error: any Error, completion: @Sendable @escaping (Result<[Flag], any Error>) -> Void) {
+        if self.defaultFlags.isEmpty {
+            completion(.failure(error))
+        } else {
+            completion(.success(self.defaultFlags))
+        }
+    }
 
     /// Check feature exists and is enabled optionally for a specific identity
     ///
@@ -268,21 +284,6 @@ public final class Flagsmith: @unchecked Sendable {
             completion(result.map(\.traits))
         }
     }
-    
-    /// Get flags for an identity and set traits in the same call
-    ///
-    /// - Parameters:
-    ///   - traits: Traits to be created or updated
-    ///   - identity: ID of the user
-    ///   - completion: Closure with Result which contains a list of Flags in case of success or Error in case of failure
-    public func getIdentityFlags(_ traits: [Trait],
-                                 forIdentity identity: String,
-                                 completion: @Sendable @escaping (Result<[Flag], any Error>) -> Void)
-    {
-        apiManager.request(.postTraits(identity: identity, traits: traits)) { (result: Result<Traits, Error>) in
-            completion(result.map(\.flags))
-        }
-    }
 
     /// Get both feature flags and user traits for the provided identity
     ///
diff --git a/FlagsmithClient/Classes/FlagsmithError.swift b/FlagsmithClient/Classes/FlagsmithError.swift
index 2500e66..41403e1 100644
--- a/FlagsmithClient/Classes/FlagsmithError.swift
+++ b/FlagsmithClient/Classes/FlagsmithError.swift
@@ -21,6 +21,8 @@ public enum FlagsmithError: LocalizedError, Sendable {
     case decoding(DecodingError)
     /// Unknown or unhandled error was encountered.
     case unhandled(any Error)
+    /// Invalid argument error
+    case invalidArgument(String)
 
     public var errorDescription: String? {
         switch self {
@@ -36,6 +38,8 @@ public enum FlagsmithError: LocalizedError, Sendable {
             return "API Response could not be decoded: \(error.localizedDescription)"
         case let .unhandled(error):
             return "An unknown or unhandled error was encountered: \(error.localizedDescription)"
+        case let .invalidArgument(error):
+            return "Invalid argument error: \(error)"
         }
     }