Add unit tests for pileup availability logic

test/python/WMComponent_t/WorkflowUpdater_t/WorkflowUpdaterPoller_t.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python
+
+import os
+import unittest
+
+from WMComponent.WorkflowUpdater.WorkflowUpdaterPoller import \
+    blockLocations, checkChanges, updateBlockInfo, writePileupJson, \
+    tarMode, findJsonSandboxFiles, extractPileupconf
+from WMCore.MicroService.Tools.Common import getMSLogger
+from WMCore.WMBase import getTestBase
+from WMQuality.Emulators.EmulatedUnitTestCase import EmulatedUnitTestCase
+
+
+class WorkflowUpdaterPollerTest(EmulatedUnitTestCase):
+    """
+    Test case for the WorkflowUpdaterPoller
+    """
+
+    def setUp(self):
+        """
+        Setup the database and logging connection. Try to create all of the
+        WMBS tables.
+        """
+        self.logger = getMSLogger(False)
+        super(WorkflowUpdaterPollerTest, self).setUp()
+
+        testData = os.path.join(getTestBase(), "WMComponent_t/WorkflowUpdater_t")
+        sandbox1 = 'SC_MultiPU_Agent227_Val_240110_215719_7133-Sandbox.tar.bz2'
+        self.sandbox1 = os.path.join(testData, sandbox1)
+        self.jdoc = {"mc": {"block1": {"FileList": [], "NumberOfEvents": 1, "PhEDExNodeNames": ['rse1']},
+                            "block2": {"FileList": [], "NumberOfEvents": 1, "PhEDExNodeNames": ['rse2']},
+                            "block3": {"FileList": [], "NumberOfEvents": 1, "PhEDExNodeNames": ['rse3']}}}
+
+    def testTarMode(self):
+        """
+        Test tarMode function
+        """
+        mode = tarMode(self.sandbox1, 'w')
+        self.assertEqual(mode, 'w:bz2')
+        mode = tarMode('file.tar', 'r')
+        self.assertEqual(mode, 'r')
+
+    def testFindJsonSandboxFiles(self):
+        """
+        Test findJsonSandboxFiles function
+        """
+        files = findJsonSandboxFiles(self.sandbox1)
+        self.assertEqual(len(files), 2)
+        self.logger.info(files)
+        expect = ['WMSandbox/GenSimFull/cmsRun1/pileupconf.json', 'WMSandbox/GenSimFull/cmsRun2/pileupconf.json']
+        self.assertEqual(files, expect)
+
+    def testExtractPileupconf(self):
+        """
+        Test extractPileupconf function
+        """
+        self.logger.info(self.sandbox1)
+        files = findJsonSandboxFiles(self.sandbox1)
+        for fname in files:
+            jdoc = extractPileupconf(self.sandbox1, fname)
+            keys = list(jdoc.keys())
+            self.assertEqual(keys, ['mc'])
+            bdict = blockLocations(jdoc)
+            blocks = bdict.keys()
+            self.logger.info("%s found %d block names", fname, len(blocks))
+            if fname == 'WMSandbox/GenSimFull/cmsRun1/pileupconf.json':
+                self.assertEqual(len(blocks), 1)
+            elif fname == 'WMSandbox/GenSimFull/cmsRun2/pileupconf.json':
+                self.assertEqual(len(blocks), 463)
+
+    def testBlockLocations(self):
+        """
+        Test blockLocations function
+        """
+        bdict = blockLocations(dict(self.jdoc))
+        self.assertEqual(list(bdict.keys()), ["block1", "block2", "block3"])
+        self.assertEqual(bdict['block1'], ['rse1'])
+        self.assertEqual(bdict['block2'], ['rse2'])
+        self.assertEqual(bdict['block3'], ['rse3'])
+
+    def testUpdataBlockInfo(self):
+        """
+        Test updateBlockInfo function
+        """
+        jdoc = dict(self.jdoc)
+        rses = ['rse1', 'rse2', 'rse3']
+        rdict = {'block1': rses}
+        bdict = {'block1': ['rse1'], 'block2': ['rse2']}
+        doc = updateBlockInfo(jdoc, rdict, bdict)
+        for rec in doc.values():
+            self.assertEqual(rec['block1']['PhEDExNodeNames'], rses)
+            # block3 should be discarded by now
+            self.assertEqual(len(rec), 2)
+            self.assertEqual(rec['block2']['PhEDExNodeNames'], ['rse2'])
+
+    def testCheckChanges(self):
+        """
+        Test checkChanges function
+        """
+        rdict = {}
+        bdict = blockLocations(dict(self.jdoc))
+        status = checkChanges(rdict, bdict)
+        self.assertEqual(status, True)
+        status = checkChanges(bdict, bdict)
+        self.assertEqual(status, False)
+
+    def testWritePileupJson(self):
+        """
+        Test writePileupJson function
+        """
+        testDoc = {"test": True, "int": 1, "list": [1, 2, 3]}
+        jdict = {}
+        tmpFile = '/tmp/WMComponent-Sandbox.tar.bz2'
+
+        # write new content to tmpFile
+        files = findJsonSandboxFiles(self.sandbox1)
+        for fname in files:
+            jdoc = extractPileupconf(self.sandbox1, fname)
+            # we skip updating part and write our test doc
+            jdict[fname] = testDoc
+        self.logger.info("Write %s to %s", jdict, tmpFile)
+        writePileupJson(self.sandbox1, jdict, tmpFile)
+
+        # now we extract from tmpFile our pileupconf.json files and make comparison of their content
+        for fname in files:
+            jdoc = extractPileupconf(tmpFile, fname)
+            self.logger.info("Extracted %s from %s (%s)", jdoc, tmpFile, fname)
+            self.assertEqual(jdoc, testDoc)
+
+        # remove tmpFile
+        os.remove(tmpFile)
+
+
+if __name__ == "__main__":
+    unittest.main()