Merge branch 'master' into tickets/DM-14563

lsst · May 25, 2018 · 3c0ebb5 · 3c0ebb5
2 parents 746ca1d + 0470d6e
commit 3c0ebb5
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 8 deletions.
diff --git a/python/lsst/daf/butler/core/quantum.py b/python/lsst/daf/butler/core/quantum.py
@@ -115,10 +115,7 @@ def addPredictedInput(self, ref):
             Reference for a Dataset to add to the Quantum's predicted inputs.
         """
         datasetTypeName = ref.datasetType.name
-        if datasetTypeName not in self._actualInputs:
-            self._predictedInputs[datasetTypeName] = [ref, ]
-        else:
-            self._predictedInputs[datasetTypeName].append(ref)
+        self._predictedInputs.setdefault(datasetTypeName, []).append(ref)
 
     def _markInputUsed(self, ref):
         """Mark an input as used.
@@ -133,10 +130,7 @@ def _markInputUsed(self, ref):
         if ref not in self._predictedInputs[datasetTypeName]:
             raise ValueError("Actual input {} was not predicted".format(ref))
         # Now insert as actual
-        if datasetTypeName not in self._actualInputs:
-            self._actualInputs[datasetTypeName] = [ref, ]
-        else:
-            self._actualInputs[datasetTypeName].append(ref)
+        self._actualInputs.setdefault(datasetTypeName, []).append(ref)
 
     def addOutput(self, ref):
         """Add an output `DatasetRef` to the `Quantum`.

diff --git a/tests/test_quantum.py b/tests/test_quantum.py
@@ -24,7 +24,9 @@
 
 import lsst.utils.tests
 
+from lsst.daf.butler.core.datasets import DatasetType, DatasetRef
 from lsst.daf.butler.core.quantum import Quantum
+from lsst.daf.butler.core.storageClass import StorageClass
 
 """Tests for Quantum.
 """
@@ -54,6 +56,46 @@ def testConstructor(self):
         self.assertEqual(quantum.endTime, endTime)
         self.assertEqual(quantum.host, host)
 
+    def testAddInputsOutputs(self):
+        """Test of addPredictedInput() method.
+        """
+        quantum = Quantum(task="some.task.object", run=None)
+
+        # start with empty
+        self.assertEqual(quantum.predictedInputs, dict())
+
+        camera = "DummyCam"
+        datasetTypeName = "test_ds"
+        storageClass = StorageClass("testref_StructuredData")
+        datasetType = DatasetType(datasetTypeName, ("camera", "visit"), storageClass)
+
+        # add one ref
+        ref = DatasetRef(datasetType, dict(camera=camera, visit=42))
+        quantum.addPredictedInput(ref)
+        self.assertIn(datasetTypeName, quantum.predictedInputs)
+        self.assertEqual(len(quantum.predictedInputs[datasetTypeName]), 1)
+        # add second ref
+        ref = DatasetRef(datasetType, dict(camera=camera, visit=43))
+        quantum.addPredictedInput(ref)
+        self.assertEqual(len(quantum.predictedInputs[datasetTypeName]), 2)
+
+        # mark last ref as actually used
+        self.assertEqual(quantum.actualInputs, dict())
+        quantum._markInputUsed(ref)
+        self.assertIn(datasetTypeName, quantum.actualInputs)
+        self.assertEqual(len(quantum.actualInputs[datasetTypeName]), 1)
+
+        # add couple of outputs too
+        self.assertEqual(quantum.outputs, dict())
+        ref = DatasetRef(datasetType, dict(camera=camera, visit=42))
+        quantum.addOutput(ref)
+        self.assertIn(datasetTypeName, quantum.outputs)
+        self.assertEqual(len(quantum.outputs[datasetTypeName]), 1)
+
+        ref = DatasetRef(datasetType, dict(camera=camera, visit=43))
+        quantum.addOutput(ref)
+        self.assertEqual(len(quantum.outputs[datasetTypeName]), 2)
+
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):
     pass