Fix incomplete map verification

enum_with_dict/enum_with_dict.py

@@ -121,12 +121,14 @@ def map(cls, key_mapping: Dict[Enum, Any]) -> Dict[Enum, Any]:
         Raises:
             KeyError: If any enum members are missing in the provided mapping.
         """
-        # Validate the new mapping
-        cls.validate_mapping_keys({**cls.to_dict(), **{member.name: value for member, value in key_mapping.items()}})
-
+
         # Map enum members to values
         mapped_values = {}
         for member in cls:
             if member in key_mapping:
                 mapped_values[member.name] = key_mapping[member]
+
+        # Validate the provided key mapping contains all enum members
+        cls.validate_mapping_keys(mapped_values)
+
         return mapped_values

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='enum_with_dict',
-    version='0.6.1',
+    version='0.6.2',
     packages=find_packages(),
     author='Jeremy Harris',
     author_email='[email protected]',

test/test_enum_with_dict.py

@@ -136,7 +136,24 @@ class TestEnum(EnumWithDict):
         }
 
         # Perform the mapping and validate
-        with self.assertRaises(AttributeError):
+        with self.assertRaises(KeyError):
+            TestEnum.map(key_mapping)
+
+    def test_incomplete_map(self):
+        """Test incomplete mapping enum members to values."""
+        class TestEnum(EnumWithDict):
+            VALUE_1 = 1
+            VALUE_2 = 2
+            VALUE_3 = 3
+
+        # Define the key mapping
+        key_mapping = {
+            TestEnum.VALUE_1: "some_new_value",
+            TestEnum.VALUE_2: "another_new_value",
+        }
+
+        # Perform the mapping and validate
+        with self.assertRaises(KeyError):
             TestEnum.map(key_mapping)
 
     def test_keys_retrieval(self):