Merge pull request #12 from christianhelle/enums

christianhelle · Jul 18, 2023 · ffbbf53 · ffbbf53
2 parents b96451f + 5535eca
commit ffbbf53
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 1 deletion.
diff --git a/src/autofaker/enums.py b/src/autofaker/enums.py
@@ -1,10 +1,11 @@
 import random
+from enum import Enum
 
 from autofaker.base import TypeDataGeneratorBase
 
 
 def is_enum(t) -> bool:
-    return t.__base__.__name__ == "Enum"
+    return isinstance(t, Enum.__class__)
 
 
 class EnumGenerator(TypeDataGeneratorBase):

diff --git a/tests/pytests/test_decorator_anonymous_enum_classes.py b/tests/pytests/test_decorator_anonymous_enum_classes.py
@@ -21,3 +21,22 @@ def test_create_enum_class_returns_not_none(sut):
 @autodata(Weekday)
 def test_create_enum_class_returns_instance(sut):
     assert isinstance(sut, Weekday)
+
+
+class BasicEnum(Enum):
+    """empty"""
+
+
+class InheritedEnum(BasicEnum):
+    VALUE_1 = 1
+    VALUE_2 = 2
+
+
+@autodata(InheritedEnum)
+def test_create_enum_class_returns_not_none(sut):
+    assert sut is not None
+
+
+@autodata(InheritedEnum)
+def test_create_enum_class_returns_instance(sut):
+    assert isinstance(sut, BasicEnum)
diff --git a/tests/test_create_anonymous_enum_classes.py b/tests/test_create_anonymous_enum_classes.py
@@ -20,3 +20,20 @@ def test_create_enum_class_returns_not_none(self):
 
     def test_create_enum_class_returns_instance(self):
         self.assertIsInstance(Autodata.create(Weekday), Weekday)
+
+
+class BasicEnum(Enum):
+    """empty"""
+
+
+class InheritedEnum(BasicEnum):
+    VALUE_1 = 1
+    VALUE_2 = 2
+
+
+class AnonymousEnumInheritanceTestCase(unittest.TestCase):
+    def test_create_enum_class_returns_not_none(self):
+        self.assertIsNotNone(Autodata.create(InheritedEnum))
+
+    def test_create_enum_class_returns_instance(self):
+        self.assertIsInstance(Autodata.create(InheritedEnum), BasicEnum)