gh-110525: Add tests for internal `set` CAPI (GH-110630)

6 files changed, 117 insertions, 3 deletions

diff --git a/Lib/test/test_capi/test_set.py b/Lib/test/test_capi/test_set.py
index e9165e7..5235f81 100644
--- a/Lib/test/test_capi/test_set.py
+++ b/Lib/test/test_capi/test_set.py
@@ -2,8 +2,9 @@ import unittest
 
 from test.support import import_helper
 
-# Skip this test if the _testcapi module isn't available.
+# Skip this test if the _testcapi or _testinternalcapi modules aren't available.
 _testcapi = import_helper.import_module('_testcapi')
+_testinternalcapi = import_helper.import_module('_testinternalcapi')
 
 class set_subclass(set):
     pass
@@ -12,13 +13,15 @@ class frozenset_subclass(frozenset):
     pass
 
 
-class TestSetCAPI(unittest.TestCase):
+class BaseSetTests:
     def assertImmutable(self, action, *args):
         self.assertRaises(SystemError, action, frozenset(), *args)
         self.assertRaises(SystemError, action, frozenset({1}), *args)
         self.assertRaises(SystemError, action, frozenset_subclass(), *args)
         self.assertRaises(SystemError, action, frozenset_subclass({1}), *args)
 
+
+class TestSetCAPI(BaseSetTests, unittest.TestCase):
     def test_set_check(self):
         check = _testcapi.set_check
         self.assertTrue(check(set()))
@@ -213,3 +216,50 @@ class TestSetCAPI(unittest.TestCase):
             clear(object())
         self.assertImmutable(clear)
         # CRASHES: clear(NULL)
+
+
+class TestInternalCAPI(BaseSetTests, unittest.TestCase):
+    def test_set_update(self):
+        update = _testinternalcapi.set_update
+        for cls in (set, set_subclass):
+            for it in ('ab', ('a', 'b'), ['a', 'b'],
+                       set('ab'), set_subclass('ab'),
+                       frozenset('ab'), frozenset_subclass('ab')):
+                with self.subTest(cls=cls, it=it):
+                    instance = cls()
+                    self.assertEqual(update(instance, it), 0)
+                    self.assertEqual(instance, {'a', 'b'})
+                    instance = cls(it)
+                    self.assertEqual(update(instance, it), 0)
+                    self.assertEqual(instance, {'a', 'b'})
+            with self.assertRaisesRegex(TypeError, 'object is not iterable'):
+                update(cls(), 1)
+            with self.assertRaisesRegex(TypeError, "unhashable type: 'dict'"):
+                update(cls(), [{}])
+        with self.assertRaises(SystemError):
+            update(object(), 'ab')
+        self.assertImmutable(update, 'ab')
+        # CRASHES: update(NULL, object())
+        # CRASHES: update(instance, NULL)
+        # CRASHES: update(NULL, NULL)
+
+    def test_set_next_entry(self):
+        set_next = _testinternalcapi.set_next_entry
+        for cls in (set, set_subclass, frozenset, frozenset_subclass):
+            with self.subTest(cls=cls):
+                instance = cls('abc')
+                pos = 0
+                items = []
+                while True:
+                    res = set_next(instance, pos)
+                    if res is None:
+                        break
+                    rc, pos, hash_, item = res
+                    items.append(item)
+                    self.assertEqual(rc, 1)
+                    self.assertIn(item, instance)
+                    self.assertEqual(hash(item), hash_)
+                self.assertEqual(items, list(instance))
+        with self.assertRaises(SystemError):
+            set_next(object(), 0)
+        # CRASHES: set_next(NULL, 0)
diff --git a/Modules/Setup.stdlib.in b/Modules/Setup.stdlib.in
index 8428142..647f442 100644
--- a/Modules/Setup.stdlib.in
+++ b/Modules/Setup.stdlib.in
@@ -158,7 +158,7 @@
 @MODULE_XXSUBTYPE_TRUE@xxsubtype xxsubtype.c
 @MODULE__XXTESTFUZZ_TRUE@_xxtestfuzz _xxtestfuzz/_xxtestfuzz.c _xxtestfuzz/fuzzer.c
 @MODULE__TESTBUFFER_TRUE@_testbuffer _testbuffer.c
-@MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c _testinternalcapi/test_lock.c _testinternalcapi/pytime.c
+@MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c _testinternalcapi/test_lock.c _testinternalcapi/pytime.c _testinternalcapi/set.c
 @MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/abstract.c _testcapi/unicode.c _testcapi/dict.c _testcapi/set.c _testcapi/getargs.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/buffer.c _testcapi/pyatomic.c _testcapi/pyos.c _testcapi/immortal.c _testcapi/heaptype_relative.c _testcapi/gc.c
 @MODULE__TESTCLINIC_TRUE@_testclinic _testclinic.c
 @MODULE__TESTCLINIC_LIMITED_TRUE@_testclinic_limited _testclinic_limited.c
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
index 05bac09..ddeb389 100644
--- a/Modules/_testinternalcapi.c
+++ b/Modules/_testinternalcapi.c
@@ -1602,6 +1602,9 @@ module_exec(PyObject *module)
     if (_PyTestInternalCapi_Init_PyTime(module) < 0) {
         return 1;
     }
+    if (_PyTestInternalCapi_Init_Set(module) < 0) {
+        return 1;
+    }
 
     if (PyModule_Add(module, "SIZEOF_PYGC_HEAD",
                         PyLong_FromSsize_t(sizeof(PyGC_Head))) < 0) {
diff --git a/Modules/_testinternalcapi/parts.h b/Modules/_testinternalcapi/parts.h
index bbb8e62..3d2774e 100644
--- a/Modules/_testinternalcapi/parts.h
+++ b/Modules/_testinternalcapi/parts.h
@@ -12,5 +12,6 @@
 
 int _PyTestInternalCapi_Init_Lock(PyObject *module);
 int _PyTestInternalCapi_Init_PyTime(PyObject *module);
+int _PyTestInternalCapi_Init_Set(PyObject *module);
 
 #endif // Py_TESTINTERNALCAPI_PARTS_H
diff --git a/Modules/_testinternalcapi/set.c b/Modules/_testinternalcapi/set.c
new file mode 100644
index 0000000..0305a78
--- /dev/null
+++ b/Modules/_testinternalcapi/set.c
@@ -0,0 +1,59 @@
+#include "parts.h"
+#include "../_testcapi/util.h"  // NULLABLE, RETURN_INT
+
+#include "pycore_setobject.h"
+
+
+static PyObject *
+set_update(PyObject *self, PyObject *args)
+{
+    PyObject *set, *iterable;
+    if (!PyArg_ParseTuple(args, "OO", &set, &iterable)) {
+        return NULL;
+    }
+    NULLABLE(set);
+    NULLABLE(iterable);
+    RETURN_INT(_PySet_Update(set, iterable));
+}
+
+static PyObject *
+set_next_entry(PyObject *self, PyObject *args)
+{
+    int rc;
+    Py_ssize_t pos;
+    Py_hash_t hash = (Py_hash_t)UNINITIALIZED_SIZE;
+    PyObject *set, *item = UNINITIALIZED_PTR;
+    if (!PyArg_ParseTuple(args, "On", &set, &pos)) {
+        return NULL;
+    }
+    NULLABLE(set);
+
+    rc = _PySet_NextEntry(set, &pos, &item, &hash);
+    if (rc == 1) {
+        return Py_BuildValue("innO", rc, pos, hash, item);
+    }
+    assert(item == UNINITIALIZED_PTR);
+    assert(hash == (Py_hash_t)UNINITIALIZED_SIZE);
+    if (rc == -1) {
+        return NULL;
+    }
+    assert(rc == 0);
+    Py_RETURN_NONE;
+}
+
+
+static PyMethodDef TestMethods[] = {
+    {"set_update", set_update, METH_VARARGS},
+    {"set_next_entry", set_next_entry, METH_VARARGS},
+
+    {NULL},
+};
+
+int
+_PyTestInternalCapi_Init_Set(PyObject *m)
+{
+    if (PyModule_AddFunctions(m, TestMethods) < 0) {
+        return -1;
+    }
+    return 0;
+}
diff --git a/PCbuild/_testinternalcapi.vcxproj b/PCbuild/_testinternalcapi.vcxproj
index fb474f0..a729ab3 100644
--- a/PCbuild/_testinternalcapi.vcxproj
+++ b/PCbuild/_testinternalcapi.vcxproj
@@ -96,6 +96,7 @@
     <ClCompile Include="..\Modules\_testinternalcapi.c" />
     <ClCompile Include="..\Modules\_testinternalcapi\pytime.c" />
     <ClCompile Include="..\Modules\_testinternalcapi\test_lock.c" />
+    <ClCompile Include="..\Modules\_testinternalcapi\set.c" />
   </ItemGroup>
   <ItemGroup>
     <ResourceCompile Include="..\PC\python_nt.rc" />