From d6695c867ad8776960c6c09615dcd52ff8161aff Mon Sep 17 00:00:00 2001
From: Philipp Moritz <pcmoritz@gmail.com>
Date: Sun, 25 Dec 2016 23:43:01 -0800
Subject: [PATCH] fix wait test (#158)

---
 src/plasma/plasma_extension.c | 24 ++++++++++++------------
 src/plasma/test/test.py       |  5 +++--
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/plasma/plasma_extension.c b/src/plasma/plasma_extension.c
index 166f21d90..c0253d0f0 100644
--- a/src/plasma/plasma_extension.c
+++ b/src/plasma/plasma_extension.c
@@ -18,7 +18,7 @@ static int PyObjectToPlasmaConnection(PyObject *object,
   }
 }
 
-static int PyObjectToUniqueID(PyObject *object, object_id *object_id) {
+static int PyStringToUniqueID(PyObject *object, object_id *object_id) {
   if (PyBytes_Check(object)) {
     memcpy(&object_id->id[0], PyBytes_AsString(object), UNIQUE_ID_SIZE);
     return 1;
@@ -61,7 +61,7 @@ PyObject *PyPlasma_create(PyObject *self, PyObject *args) {
   long long size;
   PyObject *metadata;
   if (!PyArg_ParseTuple(args, "O&O&LO", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id, &size, &metadata)) {
+                        PyStringToUniqueID, &object_id, &size, &metadata)) {
     return NULL;
   }
   if (!PyByteArray_Check(metadata)) {
@@ -89,7 +89,7 @@ PyObject *PyPlasma_hash(PyObject *self, PyObject *args) {
   plasma_connection *conn;
   object_id object_id;
   if (!PyArg_ParseTuple(args, "O&O&", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id)) {
+                        PyStringToUniqueID, &object_id)) {
     return NULL;
   }
   unsigned char digest[DIGEST_SIZE];
@@ -107,7 +107,7 @@ PyObject *PyPlasma_seal(PyObject *self, PyObject *args) {
   plasma_connection *conn;
   object_id object_id;
   if (!PyArg_ParseTuple(args, "O&O&", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id)) {
+                        PyStringToUniqueID, &object_id)) {
     return NULL;
   }
   plasma_seal(conn, object_id);
@@ -118,7 +118,7 @@ PyObject *PyPlasma_release(PyObject *self, PyObject *args) {
   plasma_connection *conn;
   object_id object_id;
   if (!PyArg_ParseTuple(args, "O&O&", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id)) {
+                        PyStringToUniqueID, &object_id)) {
     return NULL;
   }
   plasma_release(conn, object_id);
@@ -129,7 +129,7 @@ PyObject *PyPlasma_get(PyObject *self, PyObject *args) {
   plasma_connection *conn;
   object_id object_id;
   if (!PyArg_ParseTuple(args, "O&O&", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id)) {
+                        PyStringToUniqueID, &object_id)) {
     return NULL;
   }
   int64_t size;
@@ -163,7 +163,7 @@ PyObject *PyPlasma_contains(PyObject *self, PyObject *args) {
   plasma_connection *conn;
   object_id object_id;
   if (!PyArg_ParseTuple(args, "O&O&", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id)) {
+                        PyStringToUniqueID, &object_id)) {
     return NULL;
   }
   int has_object;
@@ -189,7 +189,7 @@ PyObject *PyPlasma_fetch(PyObject *self, PyObject *args) {
   Py_ssize_t n = PyList_Size(object_id_list);
   object_id *object_ids = malloc(sizeof(object_id) * n);
   for (int i = 0; i < n; ++i) {
-    PyObjectToUniqueID(PyList_GetItem(object_id_list, i), &object_ids[i]);
+    PyStringToUniqueID(PyList_GetItem(object_id_list, i), &object_ids[i]);
   }
   plasma_fetch(conn, (int) n, object_ids);
   free(object_ids);
@@ -231,8 +231,8 @@ PyObject *PyPlasma_wait(PyObject *self, PyObject *args) {
 
   object_request *object_requests = malloc(sizeof(object_request) * n);
   for (int i = 0; i < n; ++i) {
-    PyObjectToUniqueID(PyList_GetItem(object_id_list, i),
-                       &object_requests[i].object_id);
+    CHECK(PyStringToUniqueID(PyList_GetItem(object_id_list, i),
+                             &object_requests[i].object_id) == 1);
     object_requests[i].type = PLASMA_QUERY_ANYWHERE;
   }
   /* Drop the global interpreter lock while we are waiting, so other threads can
@@ -286,7 +286,7 @@ PyObject *PyPlasma_delete(PyObject *self, PyObject *args) {
   plasma_connection *conn;
   object_id object_id;
   if (!PyArg_ParseTuple(args, "O&O&", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id)) {
+                        PyStringToUniqueID, &object_id)) {
     return NULL;
   }
   plasma_delete(conn, object_id);
@@ -299,7 +299,7 @@ PyObject *PyPlasma_transfer(PyObject *self, PyObject *args) {
   const char *addr;
   int port;
   if (!PyArg_ParseTuple(args, "O&O&si", PyObjectToPlasmaConnection, &conn,
-                        PyObjectToUniqueID, &object_id, &addr, &port)) {
+                        PyStringToUniqueID, &object_id, &addr, &port)) {
     return NULL;
   }
 
diff --git a/src/plasma/test/test.py b/src/plasma/test/test.py
index 4fee0b9af..13631e611 100644
--- a/src/plasma/test/test.py
+++ b/src/plasma/test/test.py
@@ -592,8 +592,9 @@ class TestPlasmaManager(unittest.TestCase):
 
     # Make sure that wait returns when the requested number of object IDs are
     # available and does not wait for all object IDs to be available.
-    object_ids = [random_object_id() for _ in range(10)]
-    object_ids_perm = np.random.permutation(object_ids)
+    object_ids = [random_object_id() for _ in range(9)] + [20 * b'\x00']
+    object_ids_perm = object_ids[:]
+    random.shuffle(object_ids_perm)
     for i in range(10):
       if i % 2 == 0:
         create_object_with_id(self.client1, object_ids_perm[i], 2000, 2000)