From e0f74878aa557774a9e8966b2ef1fb94782d25ca Mon Sep 17 00:00:00 2001
From: davepallot <dave.pallot@uwa.edu.au>
Date: Tue, 27 Apr 2021 13:09:57 +0800
Subject: [PATCH] YAN-678: Allow uid to be the plasma object id. If uid is not
 20 bytes then create one at random.

---
 daliuge-engine/dlg/drop.py              | 11 ++++++++---
 daliuge-engine/test/apps/test_plasma.py |  8 ++++++--
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/daliuge-engine/dlg/drop.py b/daliuge-engine/dlg/drop.py
index ed6fdf164..7e1875616 100644
--- a/daliuge-engine/dlg/drop.py
+++ b/daliuge-engine/dlg/drop.py
@@ -42,6 +42,7 @@
 
 import six
 from six import BytesIO
+import numpy as np
 
 from .ddap_protocol import ExecutionMode, ChecksumTypes, AppDROPStates, \
     DROPLinkType, DROPPhases, DROPStates, DROPRel
@@ -1725,14 +1726,18 @@ class PlasmaDROP(AbstractDROP):
     plasma_link = dlg_string_param('plasma_link', '/tmp/plasma')
 
     def initialize(self, **kwargs):
-       pass
+        object_id = self.uid
+        if len(self.uid) != 20:
+            object_id = np.random.bytes(20)
+        if self.object_id is None:
+           self.object_id = object_id
 
     def getIO(self):
-        return PlasmaIO(plasma.ObjectID(self.uid), self.plasma_link)
+        return PlasmaIO(plasma.ObjectID(self.object_id), self.plasma_link)
 
     @property
     def dataURL(self):
-        return "plasma://%s" % (plasma.ObjectID(self.uid))
+        return "plasma://%s" % (self.object_id.encode('hex'))
 
 
 # Dictionary mapping 1-to-many DROPLinkType constants to the corresponding methods
diff --git a/daliuge-engine/test/apps/test_plasma.py b/daliuge-engine/test/apps/test_plasma.py
index 8b8ddfe66..e04480597 100644
--- a/daliuge-engine/test/apps/test_plasma.py
+++ b/daliuge-engine/test/apps/test_plasma.py
@@ -25,7 +25,6 @@
 
 from dlg.drop import FileDROP, PlasmaDROP
 from dlg import droputils
-import numpy as np
 
 casa_unavailable = True
 try:
@@ -65,7 +64,7 @@ def test_plasma(self):
 
         a = FileDROP('a', 'a', filepath=in_file)
         b = MSPlasmaWriter('b', 'b')
-        c = PlasmaDROP('c', np.random.bytes(20))
+        c = PlasmaDROP('c', 'c')
         d = MSPlasmaReader('d', 'd')
         e = FileDROP('e', 'e', filepath=out_file)
 
@@ -79,3 +78,8 @@ def test_plasma(self):
             a.setCompleted()
 
         self.compare_ms(in_file, out_file)
+
+        # check we can go from dataURL to plasma ID
+        client = plasma.connect("/tmp/plasma")
+        a = c.dataURL.split('//')[1].decode("hex")
+        client.get(plasma.ObjectID(a))
\ No newline at end of file