first comit

venv/lib/python3.10/site-packages/zmq/tests/test_message.py

@@ -0,0 +1,370 @@
+# Copyright (C) PyZMQ Developers
+# Distributed under the terms of the Modified BSD License.
+
+
+import copy
+import gc
+import sys
+
+try:
+    from sys import getrefcount
+except ImportError:
+    grc = None
+else:
+    grc = getrefcount
+
+import time
+
+import zmq
+from zmq.tests import PYPY, BaseZMQTestCase, SkipTest, skip_pypy
+
+# some useful constants:
+
+x = b'x'
+
+if grc:
+    rc0 = grc(x)
+    v = memoryview(x)
+    view_rc = grc(x) - rc0
+
+
+def await_gc(obj, rc):
+    """wait for refcount on an object to drop to an expected value
+
+    Necessary because of the zero-copy gc thread,
+    which can take some time to receive its DECREF message.
+    """
+    # count refs for this function
+    if sys.version_info < (3, 11):
+        my_refs = 2
+    else:
+        my_refs = 1
+    for i in range(50):
+        # rc + 2 because of the refs in this function
+        if grc(obj) <= rc + my_refs:
+            return
+        time.sleep(0.05)
+
+
+class TestFrame(BaseZMQTestCase):
+    def tearDown(self):
+        super().tearDown()
+        for i in range(3):
+            gc.collect()
+
+    @skip_pypy
+    def test_above_30(self):
+        """Message above 30 bytes are never copied by 0MQ."""
+        for i in range(5, 16):  # 32, 64,..., 65536
+            s = (2**i) * x
+            rc = grc(s)
+            m = zmq.Frame(s, copy=False)
+            assert grc(s) == rc + 2
+            del m
+            await_gc(s, rc)
+            assert grc(s) == rc
+            del s
+
+    def test_str(self):
+        """Test the str representations of the Frames."""
+        for i in range(16):
+            s = (2**i) * x
+            m = zmq.Frame(s)
+            m_str = str(m)
+            m_str_b = m_str.encode()
+            assert s == m_str_b
+
+    def test_bytes(self):
+        """Test the Frame.bytes property."""
+        for i in range(1, 16):
+            s = (2**i) * x
+            m = zmq.Frame(s)
+            b = m.bytes
+            assert s == m.bytes
+            if not PYPY:
+                # check that it copies
+                assert b is not s
+            # check that it copies only once
+            assert b is m.bytes
+
+    def test_unicode(self):
+        """Test the unicode representations of the Frames."""
+        s = 'asdf'
+        self.assertRaises(TypeError, zmq.Frame, s)
+        for i in range(16):
+            s = (2**i) * '§'
+            m = zmq.Frame(s.encode('utf8'))
+            assert s == m.bytes.decode('utf8')
+
+    def test_len(self):
+        """Test the len of the Frames."""
+        for i in range(16):
+            s = (2**i) * x
+            m = zmq.Frame(s)
+            assert len(s) == len(m)
+
+    @skip_pypy
+    def test_lifecycle1(self):
+        """Run through a ref counting cycle with a copy."""
+        for i in range(5, 16):  # 32, 64,..., 65536
+            s = (2**i) * x
+            rc = rc_0 = grc(s)
+            m = zmq.Frame(s, copy=False)
+            rc += 2
+            assert grc(s) == rc
+            m2 = copy.copy(m)
+            rc += 1
+            assert grc(s) == rc
+            # no increase in refcount for accessing buffer
+            # which references m2 directly
+            buf = m2.buffer
+            assert grc(s) == rc
+
+            assert s == str(m).encode()
+            assert s == bytes(m2)
+            assert s == m.bytes
+            assert s == bytes(buf)
+            # assert s is str(m)
+            # assert s is str(m2)
+            del m2
+            assert grc(s) == rc
+            # buf holds direct reference to m2 which holds
+            del buf
+            rc -= 1
+            assert grc(s) == rc
+            del m
+            rc -= 2
+            await_gc(s, rc)
+            assert grc(s) == rc
+            assert rc == rc_0
+            del s
+
+    @skip_pypy
+    def test_lifecycle2(self):
+        """Run through a different ref counting cycle with a copy."""
+        for i in range(5, 16):  # 32, 64,..., 65536
+            s = (2**i) * x
+            rc = rc_0 = grc(s)
+            m = zmq.Frame(s, copy=False)
+            rc += 2
+            assert grc(s) == rc
+            m2 = copy.copy(m)
+            rc += 1
+            assert grc(s) == rc
+            # no increase in refcount for accessing buffer
+            # which references m directly
+            buf = m.buffer
+            assert grc(s) == rc
+            assert s == str(m).encode()
+            assert s == bytes(m2)
+            assert s == m2.bytes
+            assert s == m.bytes
+            assert s == bytes(buf)
+            # assert s is str(m)
+            # assert s is str(m2)
+            del buf
+            assert grc(s) == rc
+            del m
+            rc -= 1
+            assert grc(s) == rc
+            del m2
+            rc -= 2
+            await_gc(s, rc)
+            assert grc(s) == rc
+            assert rc == rc_0
+            del s
+
+    def test_tracker(self):
+        m = zmq.Frame(b'asdf', copy=False, track=True)
+        assert not m.tracker.done
+        pm = zmq.MessageTracker(m)
+        assert not pm.done
+        del m
+        for i in range(3):
+            gc.collect()
+        for i in range(10):
+            if pm.done:
+                break
+            time.sleep(0.1)
+        assert pm.done
+
+    def test_no_tracker(self):
+        m = zmq.Frame(b'asdf', track=False)
+        assert m.tracker == None
+        m2 = copy.copy(m)
+        assert m2.tracker == None
+        self.assertRaises(ValueError, zmq.MessageTracker, m)
+
+    def test_multi_tracker(self):
+        m = zmq.Frame(b'asdf', copy=False, track=True)
+        m2 = zmq.Frame(b'whoda', copy=False, track=True)
+        mt = zmq.MessageTracker(m, m2)
+        assert not m.tracker.done
+        assert not mt.done
+        self.assertRaises(zmq.NotDone, mt.wait, 0.1)
+        del m
+        for i in range(3):
+            gc.collect()
+        self.assertRaises(zmq.NotDone, mt.wait, 0.1)
+        assert not mt.done
+        del m2
+        for i in range(3):
+            gc.collect()
+        assert mt.wait(0.1) is None
+        assert mt.done
+
+    def test_buffer_in(self):
+        """test using a buffer as input"""
+        ins = "§§¶•ªº˜µ¬˚…∆˙åß∂©œ∑´†≈ç√".encode()
+        zmq.Frame(memoryview(ins))
+
+    def test_bad_buffer_in(self):
+        """test using a bad object"""
+        self.assertRaises(TypeError, zmq.Frame, 5)
+        self.assertRaises(TypeError, zmq.Frame, object())
+
+    def test_buffer_out(self):
+        """receiving buffered output"""
+        ins = "§§¶•ªº˜µ¬˚…∆˙åß∂©œ∑´†≈ç√".encode()
+        m = zmq.Frame(ins)
+        outb = m.buffer
+        assert isinstance(outb, memoryview)
+        assert outb is m.buffer
+        assert m.buffer is m.buffer
+
+    def test_memoryview_shape(self):
+        """memoryview shape info"""
+        data = "§§¶•ªº˜µ¬˚…∆˙åß∂©œ∑´†≈ç√".encode()
+        n = len(data)
+        f = zmq.Frame(data)
+        view1 = f.buffer
+        assert view1.ndim == 1
+        assert view1.shape == (n,)
+        assert view1.tobytes() == data
+        view2 = memoryview(f)
+        assert view2.ndim == 1
+        assert view2.shape == (n,)
+        assert view2.tobytes() == data
+
+    def test_multisend(self):
+        """ensure that a message remains intact after multiple sends"""
+        a, b = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
+        s = b"message"
+        m = zmq.Frame(s)
+        assert s == m.bytes
+
+        a.send(m, copy=False)
+        time.sleep(0.1)
+        assert s == m.bytes
+        a.send(m, copy=False)
+        time.sleep(0.1)
+        assert s == m.bytes
+        a.send(m, copy=True)
+        time.sleep(0.1)
+        assert s == m.bytes
+        a.send(m, copy=True)
+        time.sleep(0.1)
+        assert s == m.bytes
+        for i in range(4):
+            r = b.recv()
+            assert s == r
+        assert s == m.bytes
+
+    def test_memoryview(self):
+        """test messages from memoryview"""
+        s = b'carrotjuice'
+        memoryview(s)
+        m = zmq.Frame(s)
+        buf = m.buffer
+        s2 = buf.tobytes()
+        assert s2 == s
+        assert m.bytes == s
+
+    def test_noncopying_recv(self):
+        """check for clobbering message buffers"""
+        null = b'\0' * 64
+        sa, sb = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
+        for i in range(32):
+            # try a few times
+            sb.send(null, copy=False)
+            m = sa.recv(copy=False)
+            mb = m.bytes
+            # buf = memoryview(m)
+            buf = m.buffer
+            del m
+            for i in range(5):
+                ff = b'\xff' * (40 + i * 10)
+                sb.send(ff, copy=False)
+                m2 = sa.recv(copy=False)
+                b = buf.tobytes()
+                assert b == null
+                assert mb == null
+                assert m2.bytes == ff
+                assert type(m2.bytes) is bytes
+
+    def test_noncopying_memoryview(self):
+        """test non-copying memmoryview messages"""
+        null = b'\0' * 64
+        sa, sb = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
+        for i in range(32):
+            # try a few times
+            sb.send(memoryview(null), copy=False)
+            m = sa.recv(copy=False)
+            buf = memoryview(m)
+            for i in range(5):
+                ff = b'\xff' * (40 + i * 10)
+                sb.send(memoryview(ff), copy=False)
+                m2 = sa.recv(copy=False)
+                buf2 = memoryview(m2)
+                assert buf.tobytes() == null
+                assert not buf.readonly
+                assert buf2.tobytes() == ff
+                assert not buf2.readonly
+                assert type(buf) is memoryview
+
+    def test_buffer_numpy(self):
+        """test non-copying numpy array messages"""
+        try:
+            import numpy
+            from numpy.testing import assert_array_equal
+        except ImportError:
+            raise SkipTest("requires numpy")
+        rand = numpy.random.randint
+        shapes = [rand(2, 5) for i in range(5)]
+        a, b = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
+        dtypes = [int, float, '>i4', 'B']
+        for i in range(1, len(shapes) + 1):
+            shape = shapes[:i]
+            for dt in dtypes:
+                A = numpy.empty(shape, dtype=dt)
+                a.send(A, copy=False)
+                msg = b.recv(copy=False)
+
+                B = numpy.frombuffer(msg, A.dtype).reshape(A.shape)
+                assert_array_equal(A, B)
+
+            A = numpy.empty(shape, dtype=[('a', int), ('b', float), ('c', 'a32')])
+            A['a'] = 1024
+            A['b'] = 1e9
+            A['c'] = 'hello there'
+            a.send(A, copy=False)
+            msg = b.recv(copy=False)
+
+            B = numpy.frombuffer(msg, A.dtype).reshape(A.shape)
+            assert_array_equal(A, B)
+
+    @skip_pypy
+    def test_frame_more(self):
+        """test Frame.more attribute"""
+        frame = zmq.Frame(b"hello")
+        assert not frame.more
+        sa, sb = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
+        sa.send_multipart([b'hi', b'there'])
+        frame = self.recv(sb, copy=False)
+        assert frame.more
+        if zmq.zmq_version_info()[0] >= 3 and not PYPY:
+            assert frame.get(zmq.MORE)
+        frame = self.recv(sb, copy=False)
+        assert not frame.more
+        if zmq.zmq_version_info()[0] >= 3 and not PYPY:
+            assert not frame.get(zmq.MORE)