from gi.repository import Gtk

from gaphor import UML
from gaphor.UML.classes import ClassItem, EnumerationItem
from gaphor.UML.classes.classespropertypages import (
    ClassAttributes,
    ClassEnumerationLiterals,
)


def test_attribute_editing(case):
    class_item = case.create(ClassItem, UML.Class)
    model = ClassAttributes(class_item, (str, bool, object))
    model.append([None, False, None])
    path = Gtk.TreePath.new_first()
    iter = model.get_iter(path)
    model.update(iter, col=0, value="attr")

    assert model[iter][-1] is class_item.subject.ownedAttribute[0]


def test_enumeration_editing(case):
    enum_item = case.create(EnumerationItem, UML.Enumeration)
    model = ClassEnumerationLiterals(enum_item, (str, object))
    model.append([None, None])
    path = Gtk.TreePath.new_first()
    iter = model.get_iter(path)
    model.update(iter, col=0, value="enum")

    assert model[iter][-1] is enum_item.subject.ownedLiteral[0]

from gi.repository import Gtk

from gaphor import UML
from gaphor.UML.classes import ClassItem
from gaphor.UML.classes.classespropertypages import ClassAttributes


class TestClassPropertyPages:
    def test_attribute_editing(self, case):
        class_item = case.create(ClassItem, UML.Class)
        model = ClassAttributes(class_item, (str, bool, object))
        model.append([None, False, None])
        path = Gtk.TreePath.new_first()
        iter = model.get_iter(path)
        model.update(iter, col=0, value="attr")

        assert model[iter][-1] is class_item.subject.ownedAttribute[0]

--- 
+++ 
@@ -1,17 +1,30 @@
 from gi.repository import Gtk
 
 from gaphor import UML
-from gaphor.UML.classes import ClassItem
-from gaphor.UML.classes.classespropertypages import ClassAttributes
+from gaphor.UML.classes import ClassItem, EnumerationItem
+from gaphor.UML.classes.classespropertypages import (
+    ClassAttributes,
+    ClassEnumerationLiterals,
+)
 
 
-class TestClassPropertyPages:
-    def test_attribute_editing(self, case):
-        class_item = case.create(ClassItem, UML.Class)
-        model = ClassAttributes(class_item, (str, bool, object))
-        model.append([None, False, None])
-        path = Gtk.TreePath.new_first()
-        iter = model.get_iter(path)
-        model.update(iter, col=0, value="attr")
+def test_attribute_editing(case):
+    class_item = case.create(ClassItem, UML.Class)
+    model = ClassAttributes(class_item, (str, bool, object))
+    model.append([None, False, None])
+    path = Gtk.TreePath.new_first()
+    iter = model.get_iter(path)
+    model.update(iter, col=0, value="attr")
 
-        assert model[iter][-1] is class_item.subject.ownedAttribute[0]
+    assert model[iter][-1] is class_item.subject.ownedAttribute[0]
+
+
+def test_enumeration_editing(case):
+    enum_item = case.create(EnumerationItem, UML.Enumeration)
+    model = ClassEnumerationLiterals(enum_item, (str, object))
+    model.append([None, None])
+    path = Gtk.TreePath.new_first()
+    iter = model.get_iter(path)
+    model.update(iter, col=0, value="enum")
+
+    assert model[iter][-1] is enum_item.subject.ownedLiteral[0]
Add test for enumeration editing

Signed-off-by: Dan Yeaw <2591e5f46f28d303f9dc027d475a5c60d8dea17a@yeaw.me>
