def testMongoDoc(self):
"""
Test creating a MSOutputTemplate object out of a MongoDB record
"""
msOutDoc = MSOutputTemplate(self.mongoDoc, producerDoc=False)
self.assertFalse(msOutDoc["IsRelVal"])
self.assertEqual(msOutDoc["TransferStatus"], "pending")
self.assertEqual(msOutDoc["CreationTime"],
self.mongoDoc["CreationTime"])
self.assertEqual(msOutDoc["LastUpdate"], self.mongoDoc["LastUpdate"])
self.assertEqual(msOutDoc["Campaign"], self.mongoDoc["Campaign"])
self.assertEqual(msOutDoc["OutputDatasets"],
self.mongoDoc["OutputDatasets"])
self.assertEqual(msOutDoc["RequestName"], self.mongoDoc["RequestName"])
self.assertEqual(msOutDoc["RequestType"], self.mongoDoc["RequestType"])
self.assertEqual(msOutDoc["_id"], self.mongoDoc["_id"])
self.assertEqual(len(msOutDoc["OutputMap"]), 2)
self.assertItemsEqual([
msOutDoc["OutputMap"][0]["Dataset"],
msOutDoc["OutputMap"][1]["Dataset"]
], self.mongoDoc["OutputDatasets"])
self.assertTrue(msOutDoc["OutputMap"][0]["DatasetSize"] in [123, 456])
newDoc = deepcopy(self.mongoDoc)
newDoc.update({
"IsRelVal": True,
"TransferStatus": "done",
"LastUpdate": 333
})
msOutDoc = MSOutputTemplate(newDoc, producerDoc=False)
self.assertTrue(msOutDoc["IsRelVal"])
self.assertEqual(msOutDoc["TransferStatus"], "done")
self.assertEqual(msOutDoc["LastUpdate"], 333)
def testReRecoSpec(self):
"""
Test creating a MSOutputTemplate object out of a ReReco request dictionary
"""
msOutDoc = MSOutputTemplate(self.rerecoSpec)
self.assertTrue(msOutDoc["CreationTime"] > 0)
self.assertIsNone(msOutDoc["LastUpdate"])
self.assertFalse(msOutDoc["IsRelVal"])
self.assertEqual(msOutDoc["TransferStatus"], "pending")
self.assertEqual(msOutDoc["Campaign"], self.rerecoSpec["Campaign"])
self.assertEqual(msOutDoc["OutputDatasets"],
self.rerecoSpec["OutputDatasets"])
self.assertEqual(msOutDoc["RequestName"],
self.rerecoSpec["RequestName"])
self.assertEqual(msOutDoc["RequestType"],
self.rerecoSpec["RequestType"])
self.assertEqual(msOutDoc["_id"], self.rerecoSpec["_id"])
self.assertEqual(len(msOutDoc["OutputMap"]), 2)
self.assertEqual(msOutDoc["OutputMap"][0]["DatasetSize"], 0)
self.assertItemsEqual(list(msOutDoc["OutputMap"][0]),
self.outputMapKeys)
self.assertEqual(msOutDoc["OutputMap"][0]["Campaign"],
self.rerecoSpec["Campaign"])
self.assertEqual(msOutDoc["OutputMap"][1]["Campaign"],
self.rerecoSpec["Campaign"])
self.assertItemsEqual([
msOutDoc["OutputMap"][0]["Dataset"],
msOutDoc["OutputMap"][1]["Dataset"]
], self.rerecoSpec["OutputDatasets"])
def testStepChainSpec(self):
"""
Test creating a MSOutputTemplate object out of a StepChain request dictionary
"""
msOutDoc = MSOutputTemplate(self.stepchainSpec)
self.assertTrue(msOutDoc["CreationTime"] > 0)
self.assertIsNone(msOutDoc["LastUpdate"])
self.assertFalse(msOutDoc["IsRelVal"])
self.assertEqual(msOutDoc["TransferStatus"], "pending")
self.assertEqual(msOutDoc["Campaign"], self.stepchainSpec["Campaign"])
self.assertEqual(msOutDoc["OutputDatasets"],
self.stepchainSpec["OutputDatasets"])
self.assertEqual(msOutDoc["RequestName"],
self.stepchainSpec["RequestName"])
self.assertEqual(msOutDoc["RequestType"],
self.stepchainSpec["RequestType"])
self.assertEqual(msOutDoc["_id"], self.stepchainSpec["_id"])
self.assertEqual(len(msOutDoc["OutputMap"]), 2)
self.assertItemsEqual(list(msOutDoc["OutputMap"][0]),
self.outputMapKeys)
for idx in range(2):
self.assertEqual(msOutDoc["OutputMap"][idx]["DatasetSize"], 0)
if msOutDoc["OutputMap"][idx]["Campaign"] == self.stepchainSpec[
"Step1"]["Campaign"]:
self.assertEqual(msOutDoc["OutputMap"][idx]["Dataset"],
"output-dataset-1")
else:
self.assertEqual(msOutDoc["OutputMap"][idx]["Campaign"],
self.stepchainSpec["Step2"]["Campaign"])
self.assertEqual(msOutDoc["OutputMap"][idx]["Dataset"],
"output-dataset-2")
def docReadfromMongo(self, mQueryDict, dbColl, setTaken=False):
"""
Reads a single Document from MongoDB and if setTaken flag is on then
Sets the relevant flags (isTaken, isTakenBy) in the document at MongoDB
"""
# NOTE:
# In case the current query returns an empty document from MongoDB
# (eg. all workflows have been processed) the MSOutputTemplate
# will throw an error. We should catch this one here and interrupt
# the pipeLine traversal, otherwise an error either here or in one of the
# following stages will most probably occur and the whole run will be broken.
if setTaken:
lastUpdate = int(time())
retrString = self.currThreadIdent
mongoDoc = dbColl.find_one_and_update(
mQueryDict, {
'$set': {
'isTaken': True,
'isTakenBy': retrString,
'lastUpdate': lastUpdate
}
},
return_document=ReturnDocument.AFTER)
else:
mongoDoc = dbColl.find_one(mQueryDict)
if mongoDoc:
try:
msOutDoc = MSOutputTemplate(mongoDoc)
except Exception as ex:
msg = "Unable to create msOutDoc from %s." % mongoDoc
self.logger.warning(msg)
raise ex
# NOTE:
# Here if we do not update the isTaken flag in MongoDB back
# to False, the document will not be released in MongoDb and
# will stay locked. If we are ending up here it means for some
# reason we have a malformed document in MongoDB. We should make
# a design choice - should we release the document or should
# we leave it locked for further investigations or maybe
# mark it with another flag eg. 'isMalformed': True
else:
raise EmptyResultError
return msOutDoc
def docTransformer(self, doc):
"""
A function used to transform a request record from reqmgr2 to a document
suitable for uploading to Mongodb
"""
# Solution 1: Destructive function - to force clear of the the externally
# referenced object and to return a new one (current solution)
# NOTE: Leaves an empty dictionary behind (the clear method just
# clears all the keys of the dict, but does not delete it)
# Solution 2: To work in place (will keep the dynamic structure of the passed dict)
# Solution 3: To have 2 object buffers for the two diff types outside the function
try:
msOutDoc = MSOutputTemplate(doc)
doc.clear()
except Exception as ex:
msg = "ERR: Unable to create MSOutputTemplate for document: \n%s\n" % pformat(
doc)
msg += "ERR: %s" % str(ex)
self.logger.exception(msg)
raise ex
return msOutDoc
def testTaskChainSpec(self):
"""
Test creating a MSOutputTemplate object out of a TaskChain request dictionary
"""
msOutDoc = MSOutputTemplate(self.taskchainSpec)
self.assertTrue(msOutDoc["CreationTime"] > 0)
self.assertIsNone(msOutDoc["LastUpdate"])
self.assertFalse(msOutDoc["IsRelVal"])
self.assertEqual(msOutDoc["TransferStatus"], "pending")
self.assertEqual(msOutDoc["Campaign"], self.taskchainSpec["Campaign"])
self.assertEqual(msOutDoc["OutputDatasets"],
self.taskchainSpec["OutputDatasets"])
self.assertEqual(msOutDoc["RequestName"],
self.taskchainSpec["RequestName"])
self.assertEqual(msOutDoc["RequestType"],
self.taskchainSpec["RequestType"])
self.assertEqual(msOutDoc["_id"], self.taskchainSpec["_id"])
self.assertEqual(len(msOutDoc["OutputMap"]), 2)
self.assertItemsEqual(msOutDoc["OutputMap"][0].viewkeys(),
self.outputMapKeys)
for idx in range(2):
self.assertEqual(msOutDoc["OutputMap"][idx]["DatasetSize"], 0)
if msOutDoc["OutputMap"][idx]["Campaign"] == self.taskchainSpec[
"Task1"]["Campaign"]:
self.assertEqual(msOutDoc["OutputMap"][idx]["Dataset"],
"output-dataset-1")
else:
self.assertEqual(msOutDoc["OutputMap"][idx]["Campaign"],
self.taskchainSpec["Task2"]["Campaign"])
self.assertEqual(msOutDoc["OutputMap"][idx]["Dataset"],
"output-dataset-2")
self.assertEqual(msOutDoc["OutputMap"][idx]["Copies"], 1)
self.assertEqual(msOutDoc["OutputMap"][idx]["DiskDestination"], "")
self.assertEqual(msOutDoc["OutputMap"][idx]["DiskRuleID"], "")
self.assertEqual(msOutDoc["OutputMap"][idx]["TapeDestination"], "")
self.assertEqual(msOutDoc["OutputMap"][idx]["TapeRuleID"], "")
def testSetters(self):
"""
Test the MSOutputTemplate setter methods
"""
msOutDoc = MSOutputTemplate(self.mongoDoc)
self.assertFalse(msOutDoc["IsRelVal"])
msOutDoc._setRelVal({"SubRequestType": "RelVal"})
msOutDoc.setKey("IsRelVal", True)
self.assertTrue(msOutDoc["IsRelVal"])
with self.assertRaises(KeyError):
msOutDoc.setKey("alan", True)
msOutDoc.updateDoc({"IsRelVal": False, "LastUpdate": 444})
self.assertFalse(msOutDoc["IsRelVal"])
self.assertEqual(msOutDoc["LastUpdate"], 444)
msOutDoc.updateTime()
self.assertTrue(msOutDoc["LastUpdate"] > 444)
with self.assertRaises(RuntimeError):
msOutDoc.setTransferStatus("bad_status")
self.assertEqual(msOutDoc["TransferStatus"], "pending")
msOutDoc.setTransferStatus("done")
self.assertEqual(msOutDoc["TransferStatus"], "done")