Fix argument handling of is/set_non_blocking()

PyGreSQL · Mar 28, 2021 · 16804ee · 16804ee
1 parent 6969e74
commit 16804ee
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 20 deletions.
diff --git a/pgconn.c b/pgconn.c
@@ -641,7 +641,9 @@ conn_set_non_blocking(connObject *self, PyObject *args)
     }
 
     if (!PyArg_ParseTuple(args, "i", &non_blocking)) {
-        PyErr_SetString(PyExc_TypeError, "setnonblocking(tf), with boolean.");
+        PyErr_SetString(
+            PyExc_TypeError,
+            "set_non_blocking() expects a boolean value as argument");
         return NULL;
     }
 
@@ -658,7 +660,7 @@ static char conn_is_non_blocking__doc__[] =
 "is_non_blocking() -- report the blocking status of the connection";
 
 static PyObject *
-conn_is_non_blocking(connObject *self, PyObject *args)
+conn_is_non_blocking(connObject *self, PyObject *noargs)
 {
     int rc;
 
@@ -667,19 +669,13 @@ conn_is_non_blocking(connObject *self, PyObject *args)
         return NULL;
     }
 
-    if (!PyArg_ParseTuple(args, "")) {
-        PyErr_SetString(PyExc_TypeError,
-            "method is_non_blocking() takes no parameters");
-        return NULL;
-    }
-
     rc = PQisnonblocking(self->cnx);
     if (rc < 0) {
         PyErr_SetString(PyExc_IOError, PQerrorMessage(self->cnx));
         return NULL;
     }
 
-    return PyBool_FromLong(rc);
+    return PyBool_FromLong((long)rc);
 }
 #endif /* DIRECT_ACCESS */
 
@@ -1426,7 +1422,7 @@ static char conn_poll__doc__[] =
 "poll() -- Completes an asynchronous connection";
 
 static PyObject *
-conn_poll(connObject *self, PyObject *args)
+conn_poll(connObject *self, PyObject *noargs)
 {
     int rc;
 
@@ -1435,13 +1431,6 @@ conn_poll(connObject *self, PyObject *args)
         return NULL;
     }
 
-    /* check args */
-    if (!PyArg_ParseTuple(args, "")) {
-        PyErr_SetString(PyExc_TypeError,
-            "method poll() takes no parameters");
-        return NULL;
-    }
-
     Py_BEGIN_ALLOW_THREADS
     rc = PQconnectPoll(self->cnx);
     Py_END_ALLOW_THREADS
@@ -1612,7 +1601,7 @@ static struct PyMethodDef conn_methods[] = {
     {"describe_prepared", (PyCFunction) conn_describe_prepared,
         METH_VARARGS, conn_describe_prepared__doc__},
     {"poll", (PyCFunction) conn_poll,
-        METH_VARARGS, conn_poll__doc__},
+        METH_NOARGS, conn_poll__doc__},
     {"reset", (PyCFunction) conn_reset,
         METH_NOARGS, conn_reset__doc__},
     {"cancel", (PyCFunction) conn_cancel,
@@ -1659,7 +1648,7 @@ static struct PyMethodDef conn_methods[] = {
     {"endcopy", (PyCFunction) conn_endcopy,
         METH_NOARGS, conn_endcopy__doc__},
     {"set_non_blocking", (PyCFunction) conn_set_non_blocking,
-        METH_O, conn_set_non_blocking__doc__},
+        METH_VARARGS, conn_set_non_blocking__doc__},
     {"is_non_blocking", (PyCFunction) conn_is_non_blocking,
         METH_NOARGS, conn_is_non_blocking__doc__},
 #endif /* DIRECT_ACCESS */

diff --git a/tests/test_classic_connection.py b/tests/test_classic_connection.py
@@ -84,22 +84,35 @@ class TestCanConnect(unittest.TestCase):
     def testCanConnect(self):
         try:
             connection = connect()
+            rc = connection.poll()
         except pg.Error as error:
             self.fail('Cannot connect to database %s:\n%s' % (dbname, error))
+        self.assertEqual(rc, pg.POLLING_OK)
+        self.assertIs(connection.is_non_blocking(), False)
+        connection.set_non_blocking(True)
+        self.assertIs(connection.is_non_blocking(), True)
+        connection.set_non_blocking(False)
+        self.assertIs(connection.is_non_blocking(), False)
         try:
             connection.close()
         except pg.Error:
             self.fail('Cannot close the database connection')
 
     def testCanConnectNoWait(self):
         try:
-            connection = connect()
+            connection = connect_nowait()
             rc = connection.poll()
+            self.assertEqual(rc, pg.POLLING_READING)
             while rc not in (pg.POLLING_OK, pg.POLLING_FAILED):
                 rc = connection.poll()
         except pg.Error as error:
             self.fail('Cannot connect to database %s:\n%s' % (dbname, error))
         self.assertEqual(rc, pg.POLLING_OK)
+        self.assertIs(connection.is_non_blocking(), False)
+        connection.set_non_blocking(True)
+        self.assertIs(connection.is_non_blocking(), True)
+        connection.set_non_blocking(False)
+        self.assertIs(connection.is_non_blocking(), False)
         try:
             connection.close()
         except pg.Error: