From a520ea249e077da533751b9f75f284c64cdbf874 Mon Sep 17 00:00:00 2001
From: evilchili <evilchili@gmail.com>
Date: Sat, 20 Apr 2024 23:27:47 -0700
Subject: [PATCH] adding tests for ancestries

---
 src/ttfrog/db/base.py             |  3 +++
 src/ttfrog/db/schema/character.py | 18 +++++++++++---
 test/conftest.py                  |  3 ++-
 test/test_schema.py               | 39 ++++++++++++++++++++++++++++++-
 4 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/src/ttfrog/db/base.py b/src/ttfrog/db/base.py
index 965828b..c15ced3 100644
--- a/src/ttfrog/db/base.py
+++ b/src/ttfrog/db/base.py
@@ -113,5 +113,8 @@ CREATURE_TYPES = [
     "plant",
     "undead",
 ]
+SIZES = ["Tiny", "Small", "Medium", "Large", "Huge", "Gargantuan"]
+
 CreatureTypesEnum = EnumField("CreatureTypesEnum", ((k, k) for k in CREATURE_TYPES))
 StatsEnum = EnumField("StatsEnum", ((k, k) for k in STATS))
+SizesEnum = EnumField("SizesEnum", ((k, k) for k in SIZES))
diff --git a/src/ttfrog/db/schema/character.py b/src/ttfrog/db/schema/character.py
index 972e290..5c04276 100644
--- a/src/ttfrog/db/schema/character.py
+++ b/src/ttfrog/db/schema/character.py
@@ -2,7 +2,7 @@ from sqlalchemy import Column, Enum, ForeignKey, Integer, String, Text, UniqueCo
 from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.orm import relationship
 
-from ttfrog.db.base import BaseObject, CreatureTypesEnum, SavingThrowsMixin, SkillsMixin, SlugMixin
+from ttfrog.db.base import BaseObject, CreatureTypesEnum, SavingThrowsMixin, SizesEnum, SkillsMixin, SlugMixin
 
 __all__ = [
     "Ancestry",
@@ -44,9 +44,21 @@ class Ancestry(BaseObject):
     __tablename__ = "ancestry"
     id = Column(Integer, primary_key=True, autoincrement=True)
     name = Column(String, index=True, unique=True)
-    creature_type = Column(Enum(CreatureTypesEnum))
+    creature_type = Column(Enum(CreatureTypesEnum), nullable=False, default="humanoid")
+    size = Column(Enum(SizesEnum), nullable=False, default="Medium")
+    speed = Column(Integer, nullable=False, default=30, info={"min": 0, "max": 99})
     _traits = relationship("AncestryTraitMap", lazy="immediate")
 
+    @property
+    def traits(self):
+        return [mapping.trait for mapping in self._traits]
+
+    def add_trait(self, trait, level=1):
+        if trait not in self.traits:
+            self._traits.append(AncestryTraitMap(ancestry_id=self.id, ancestry_trait_id=trait.id, level=level))
+            return True
+        return False
+
     def __repr__(self):
         return self.name
 
@@ -133,7 +145,7 @@ class Character(BaseObject, SlugMixin, SavingThrowsMixin, SkillsMixin):
 
     @property
     def traits(self):
-        return [mapping.trait for mapping in self.ancestry._traits]
+        return self.ancestry.traits
 
     @property
     def level(self):
diff --git a/test/conftest.py b/test/conftest.py
index d1e3412..eb8cbcc 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -26,7 +26,8 @@ def db(monkeypatch):
     monkeypatch.setenv("DATABASE_URL", "sqlite:///:memory:")
     monkeypatch.setenv("DEBUG", "1")
     _db.init()
-    return _db
+    yield _db
+    _db.metadata.drop_all(bind=_db.engine)
 
 
 @pytest.fixture
diff --git a/test/test_schema.py b/test/test_schema.py
index 35454d8..a228373 100644
--- a/test/test_schema.py
+++ b/test/test_schema.py
@@ -26,6 +26,16 @@ def test_manage_character(db, classes_factory, ancestries_factory):
         assert char.ancestry.name == "tiefling"
         assert darkvision in char.traits
 
+        # switch ancestry to dragonborn and assert darkvision persists
+        char.ancestry = ancestries["dragonborn"]
+        db.add_or_update(char)
+        assert darkvision in char.traits
+
+        # switch ancestry to human and assert darkvision is removed
+        char.ancestry = ancestries["human"]
+        db.add_or_update(char)
+        assert darkvision not in char.traits
+
         # assign a class and level
         char.add_class(classes["fighter"], level=1)
         db.add_or_update(char)
@@ -34,7 +44,8 @@ def test_manage_character(db, classes_factory, ancestries_factory):
         assert char.class_attributes == {}
 
         # 'fighting style' is available, but not at this level
-        fighting_style = char.classes["fighter"].attributes_by_level[2]["Fighting Style"]
+        fighter = classes["fighter"]
+        fighting_style = fighter.attributes_by_level[2]["Fighting Style"]
         assert char.add_class_attribute(fighting_style, fighting_style.options[0]) is False
         db.add_or_update(char)
         assert char.class_attributes == {}
@@ -71,3 +82,29 @@ def test_manage_character(db, classes_factory, ancestries_factory):
         dump = json.loads(db.dump())
         assert dump["class_map"] == []
         assert dump["character_class_attribute_map"] == []
+
+
+def test_ancestries(db):
+    with db.transaction():
+        # create the Pygmy Orc ancestry
+        porc = schema.Ancestry(
+            name="Pygmy Orc",
+            size="Small",
+            speed=25,
+        )
+        db.add_or_update(porc)
+        assert porc.name == "Pygmy Orc"
+        assert porc.creature_type == "humanoid"
+        assert porc.size == "Small"
+        assert porc.speed == 25
+
+        # create the Relentless Endurance trait and add it to the Orc
+        endurance = schema.AncestryTrait(name="Relentless Endurance")
+        db.add_or_update(endurance)
+        porc.add_trait(endurance, level=1)
+        db.add_or_update(porc)
+        assert endurance in porc.traits
+
+        # now create an orc character and assert it gets Relentless Endurance
+        grognak = schema.Character(name="Grognak the Mighty", ancestry=porc)
+        assert endurance in grognak.traits