From 950dd4751c8b5d497e13d3c9770fdc1a926e137e Mon Sep 17 00:00:00 2001
From: Erfan Nourbakhsh <erfan@princeton.edu>
Date: Tue, 10 Dec 2024 23:21:38 -0500
Subject: [PATCH] Add unit tests for keyCheck in DictField and ConfigDictField

---
 tests/test_configDictField.py | 26 ++++++++++++++++++++++++++
 tests/test_dictField.py       | 30 +++++++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/tests/test_configDictField.py b/tests/test_configDictField.py
index 7474519..d723645 100644
--- a/tests/test_configDictField.py
+++ b/tests/test_configDictField.py
@@ -54,6 +54,14 @@ class Config3(pexConfig.Config):
     field1 = pexConfig.ConfigDictField(keytype=str, itemtype=pexConfig.Config, default={}, doc="doc")
 
 
+class Config4(pexConfig.Config):
+    """Fourth test config."""
+
+    field1 = pexConfig.ConfigDictField(
+        keytype=str, itemtype=pexConfig.Config, default={}, doc="doc", keyCheck=lambda k: k.islower()
+    )
+
+
 class ConfigDictFieldTest(unittest.TestCase):
     """Test of ConfigDictField."""
 
@@ -78,6 +86,16 @@ class BadItemtype(pexConfig.Config):
         else:
             raise SyntaxError("Unsupported itemtypes should not be allowed")
 
+        try:
+
+            class BadKeyCheck(pexConfig.Config):
+                d = pexConfig.ConfigDictField("...", keytype=str, itemtype=Config1, keyCheck=4)
+
+        except Exception:
+            pass
+        else:
+            raise SyntaxError("Non-callable keyCheck should not be allowed")
+
         try:
 
             class BadItemCheck(pexConfig.Config):
@@ -116,6 +134,14 @@ def testValidate(self):
         c.d1["a"].f = 5
         c.validate()
 
+    def testKeyCheckValidation(self):
+        c = Config4()
+        c.field1["lower"] = pexConfig.Config()
+        with self.assertRaises(pexConfig.FieldValidationError, msg="Key check should fail"):
+            c.field1["UPPER"] = pexConfig.Config()
+            # No need for c.validate() here, as the exception for key check is
+            # raised by the assignment.
+
     def testInPlaceModification(self):
         c = Config2(d1={})
         self.assertRaises(pexConfig.FieldValidationError, c.d1.__setitem__, 1, 0)
diff --git a/tests/test_dictField.py b/tests/test_dictField.py
index bd4667d..8672ad9 100644
--- a/tests/test_dictField.py
+++ b/tests/test_dictField.py
@@ -38,6 +38,8 @@ class Config1(pexConfig.Config):
     d2 = pexConfig.DictField("d2", keytype=str, itemtype=str, default=None)
     d3 = pexConfig.DictField("d3", keytype=float, itemtype=float, optional=True, itemCheck=lambda x: x > 0)
     d4 = pexConfig.DictField("d4", keytype=str, itemtype=None, default={})
+    d5 = pexConfig.DictField[str, float]("d5", default={}, keyCheck=lambda k: k not in ["k1", "k2"])
+    d6 = pexConfig.DictField[int, str]("d6", default={-2: "v1", 4: "v2"}, keyCheck=lambda k: k % 2 == 0)
 
 
 class DictFieldTest(unittest.TestCase):
@@ -66,7 +68,7 @@ class BadItemtype(pexConfig.Config):
 
         try:
 
-            class BadItemCheck(pexConfig.Config):
+            class BadKeyCheck(pexConfig.Config):
                 d = pexConfig.DictField("...", keytype=int, itemtype=int, itemCheck=4)
 
         except Exception:
@@ -74,6 +76,16 @@ class BadItemCheck(pexConfig.Config):
         else:
             raise SyntaxError("Non-callable itemCheck DictFields should not be allowed")
 
+        try:
+
+            class BadItemCheck(pexConfig.Config):
+                d = pexConfig.DictField("...", keytype=str, itemtype=float, keyCheck=4)
+
+        except Exception:
+            pass
+        else:
+            raise SyntaxError("Non-callable keyCheck DictFields should not be allowed")
+
         try:
 
             class BadDictCheck(pexConfig.Config):
@@ -139,6 +151,22 @@ def testValidate(self):
         c.d2 = {"a": "b"}
         c.validate()
 
+    def testKeyCheckValidation(self):
+        c = Config1()
+        c.d5 = {"k3": -1, "k4": 0.25}
+        c.d6 = {6: "v3"}
+
+        with self.assertRaises(
+            pexConfig.FieldValidationError,
+            msg="Key check must reject dictionary assignment with invalid keys",
+        ):
+            c.d5 = {"k1": 1.5, "k2": 2.0}
+
+        with self.assertRaises(
+            pexConfig.FieldValidationError, msg="Key check must reject invalid key addition",
+        ):
+            c.d6[3] = "v4"
+
     def testInPlaceModification(self):
         c = Config1()
         self.assertRaises(pexConfig.FieldValidationError, c.d1.__setitem__, 2, 0)