From 34ec90067893af7ce38ddf40ce582e418fdde380 Mon Sep 17 00:00:00 2001
From: Tal Borenstein <tal@keephq.dev>
Date: Thu, 21 Nov 2024 20:08:50 +0200
Subject: [PATCH] fix: wip

---
 keep/api/api.py                                |  1 +
 keep/identitymanager/identitymanagerfactory.py | 17 +++++++++--------
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/keep/api/api.py b/keep/api/api.py
index 174c7e6ff..3c2ed422b 100644
--- a/keep/api/api.py
+++ b/keep/api/api.py
@@ -15,6 +15,7 @@
 
 import keep.api.logging
 import keep.api.observability
+import keep.api.utils.import_ee
 from keep.api.arq_worker import get_arq_worker
 from keep.api.consts import (
     KEEP_ARQ_QUEUE_BASIC,
diff --git a/keep/identitymanager/identitymanagerfactory.py b/keep/identitymanager/identitymanagerfactory.py
index acd2387dd..04077da4a 100644
--- a/keep/identitymanager/identitymanagerfactory.py
+++ b/keep/identitymanager/identitymanagerfactory.py
@@ -105,13 +105,14 @@ def _load_manager(manager_type: str, manager_class: str, *args, **kwargs):
                 )
             # look for the module in ee
             except ModuleNotFoundError:
-                module = importlib.import_module(
-                    f"ee.identitymanager.identity_managers.{manager_type}.{manager_type}_{manager_class}"
-                )
-            except ModuleNotFoundError:
-                raise NotImplementedError(
-                    f"{manager_class.__name__} for {manager_type} not implemented"
-                )
+                try:
+                    module = importlib.import_module(
+                        f"ee.identitymanager.identity_managers.{manager_type}.{manager_type}_{manager_class}"
+                    )
+                except ModuleNotFoundError:
+                    raise NotImplementedError(
+                        f"{manager_class} for {manager_type} not implemented"
+                    )
             # look for the class that contains the manager_class in its name
             for _attr in dir(module):
                 if manager_class in _attr.lower() and "base" not in _attr.lower():
@@ -121,7 +122,7 @@ def _load_manager(manager_type: str, manager_class: str, *args, **kwargs):
             return manager_class(*args, **kwargs)
         except (ImportError, AttributeError):
             raise NotImplementedError(
-                f"{manager_class.__name__} for {manager_type} not implemented"
+                f"{manager_class} for {manager_type} not implemented"
             )
 
     @staticmethod