Coverage for tests/test_cmdLineFwk.py : 22%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of ctrl_mpexec.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22"""Simple unit test for cmdLineFwk module.
23"""
25import argparse
26import contextlib
27import logging
28import os
29import pickle
30import tempfile
31import unittest
33from lsst.ctrl.mpexec.cmdLineFwk import CmdLineFwk
34from lsst.ctrl.mpexec.cmdLineParser import (_ACTION_ADD_TASK, _ACTION_CONFIG,
35 _ACTION_CONFIG_FILE, _ACTION_ADD_INSTRUMENT)
36from lsst.daf.butler import Config, Quantum, Registry
37from lsst.daf.butler.registry import RegistryConfig
38from lsst.obs.base import Instrument
39import lsst.pex.config as pexConfig
40from lsst.pipe.base import (Pipeline, PipelineTask, PipelineTaskConfig,
41 QuantumGraph, QuantumGraphTaskNodes,
42 TaskDef, TaskFactory, PipelineTaskConnections)
43import lsst.pipe.base.connectionTypes as cT
44import lsst.utils.tests
45from testUtil import (AddTask, AddTaskFactoryMock, makeSimpleQGraph)
48logging.basicConfig(level=logging.INFO)
50# Have to monkey-patch Instrument.fromName() to not retrieve non-existing
51# instrument from registry, these tests can run fine without actual instrument
52# and implementing full mock for Instrument is too complicated.
53Instrument.fromName = lambda name, reg: None 53 ↛ exitline 53 didn't run the lambda on line 53
56@contextlib.contextmanager
57def makeTmpFile(contents=None):
58 """Context manager for generating temporary file name.
60 Temporary file is deleted on exiting context.
62 Parameters
63 ----------
64 contents : `bytes`
65 Data to write into a file.
66 """
67 fd, tmpname = tempfile.mkstemp()
68 if contents:
69 os.write(fd, contents)
70 os.close(fd)
71 yield tmpname
72 with contextlib.suppress(OSError):
73 os.remove(tmpname)
76@contextlib.contextmanager
77def makeSQLiteRegistry():
78 """Context manager to create new empty registry database.
80 Yields
81 ------
82 config : `RegistryConfig`
83 Registry configuration for initialized registry database.
84 """
85 with makeTmpFile() as filename:
86 uri = f"sqlite:///{filename}"
87 config = RegistryConfig()
88 config["db"] = uri
89 Registry.fromConfig(config, create=True)
90 yield config
93class SimpleConnections(PipelineTaskConnections, dimensions=(),
94 defaultTemplates={"template": "simple"}):
95 schema = cT.InitInput(doc="Schema",
96 name="{template}schema",
97 storageClass="SourceCatalog")
100class SimpleConfig(PipelineTaskConfig, pipelineConnections=SimpleConnections):
101 field = pexConfig.Field(dtype=str, doc="arbitrary string")
103 def setDefaults(self):
104 PipelineTaskConfig.setDefaults(self)
107class TaskOne(PipelineTask):
108 ConfigClass = SimpleConfig
109 _DefaultName = "taskOne"
112class TaskTwo(PipelineTask):
113 ConfigClass = SimpleConfig
114 _DefaultName = "taskTwo"
117class TaskFactoryMock(TaskFactory):
118 def loadTaskClass(self, taskName):
119 if taskName == "TaskOne":
120 return TaskOne, "TaskOne"
121 elif taskName == "TaskTwo":
122 return TaskTwo, "TaskTwo"
124 def makeTask(self, taskClass, config, overrides, butler):
125 if config is None:
126 config = taskClass.ConfigClass()
127 if overrides:
128 overrides.applyTo(config)
129 return taskClass(config=config, butler=butler)
132def _makeArgs(pipeline=None, qgraph=None, pipeline_actions=(), order_pipeline=False,
133 save_pipeline="", save_qgraph="", save_single_quanta="",
134 pipeline_dot="", qgraph_dot="", registryConfig=None):
135 """Return parsed command line arguments.
137 Parameters
138 ----------
139 pipeline : `str`, optional
140 Name of the YAML file with pipeline.
141 qgraph : `str`, optional
142 Name of the pickle file with QGraph.
143 pipeline_actions : itrable of `cmdLinePArser._PipelineAction`, optional
144 order_pipeline : `bool`
145 save_pipeline : `str`
146 Name of the YAML file to store pipeline.
147 save_qgraph : `str`
148 Name of the pickle file to store QGraph.
149 save_single_quanta : `str`
150 Name of the pickle file pattern to store individual QGraph.
151 pipeline_dot : `str`
152 Name of the DOT file to write pipeline graph.
153 qgraph_dot : `str`
154 Name of the DOT file to write QGraph representation.
155 """
156 args = argparse.Namespace()
157 args.butler_config = Config()
158 if registryConfig:
159 args.butler_config["registry"] = registryConfig
160 # The default datastore has a relocatable root, so we need to specify
161 # some root here for it to use
162 args.butler_config.configFile = "."
163 args.pipeline = pipeline
164 args.qgraph = qgraph
165 args.pipeline_actions = pipeline_actions
166 args.order_pipeline = order_pipeline
167 args.save_pipeline = save_pipeline
168 args.save_qgraph = save_qgraph
169 args.save_single_quanta = save_single_quanta
170 args.pipeline_dot = pipeline_dot
171 args.qgraph_dot = qgraph_dot
172 args.input = ""
173 args.output = None
174 args.output_run = None
175 args.extend_run = False
176 args.replace_run = False
177 args.prune_replaced = False
178 args.register_dataset_types = False
179 args.skip_init_writes = False
180 args.no_versions = False
181 args.skip_existing = False
182 args.init_only = False
183 args.processes = 1
184 args.profile = None
185 args.enableLsstDebug = False
186 args.graph_fixup = None
187 return args
190def _makeQGraph():
191 """Make a trivial QuantumGraph with one quantum.
193 The only thing that we need to do with this quantum graph is to pickle
194 it, the quanta in this graph are not usable for anything else.
196 Returns
197 -------
198 qgraph : `~lsst.pipe.base.QuantumGraph`
199 """
200 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig())
201 quanta = [Quantum()]
202 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={})
203 qgraph = QuantumGraph([taskNodes])
204 return qgraph
207class CmdLineFwkTestCase(unittest.TestCase):
208 """A test case for CmdLineFwk
209 """
211 def testMakePipeline(self):
212 """Tests for CmdLineFwk.makePipeline method
213 """
214 fwk = CmdLineFwk()
216 # make empty pipeline
217 args = _makeArgs()
218 pipeline = fwk.makePipeline(args)
219 self.assertIsInstance(pipeline, Pipeline)
220 self.assertEqual(len(pipeline), 0)
222 # few tests with serialization
223 with makeTmpFile() as tmpname:
224 # make empty pipeline and store it in a file
225 args = _makeArgs(save_pipeline=tmpname)
226 pipeline = fwk.makePipeline(args)
227 self.assertIsInstance(pipeline, Pipeline)
229 # read pipeline from a file
230 args = _makeArgs(pipeline=tmpname)
231 pipeline = fwk.makePipeline(args)
232 self.assertIsInstance(pipeline, Pipeline)
233 self.assertEqual(len(pipeline), 0)
235 # single task pipeline
236 actions = [
237 _ACTION_ADD_TASK("TaskOne:task1")
238 ]
239 args = _makeArgs(pipeline_actions=actions)
240 pipeline = fwk.makePipeline(args)
241 self.assertIsInstance(pipeline, Pipeline)
242 self.assertEqual(len(pipeline), 1)
244 # many task pipeline
245 actions = [
246 _ACTION_ADD_TASK("TaskOne:task1a"),
247 _ACTION_ADD_TASK("TaskTwo:task2"),
248 _ACTION_ADD_TASK("TaskOne:task1b")
249 ]
250 args = _makeArgs(pipeline_actions=actions)
251 pipeline = fwk.makePipeline(args)
252 self.assertIsInstance(pipeline, Pipeline)
253 self.assertEqual(len(pipeline), 3)
255 # single task pipeline with config overrides, cannot use TaskOne, need
256 # something that can be imported with `doImport()`
257 actions = [
258 _ACTION_ADD_TASK("testUtil.AddTask:task"),
259 _ACTION_CONFIG("task:addend=100")
260 ]
261 args = _makeArgs(pipeline_actions=actions)
262 pipeline = fwk.makePipeline(args)
263 taskDefs = list(pipeline.toExpandedPipeline())
264 self.assertEqual(len(taskDefs), 1)
265 self.assertEqual(taskDefs[0].config.addend, 100)
267 overrides = b"config.addend = 1000\n"
268 with makeTmpFile(overrides) as tmpname:
269 actions = [
270 _ACTION_ADD_TASK("testUtil.AddTask:task"),
271 _ACTION_CONFIG_FILE("task:" + tmpname)
272 ]
273 args = _makeArgs(pipeline_actions=actions)
274 pipeline = fwk.makePipeline(args)
275 taskDefs = list(pipeline.toExpandedPipeline())
276 self.assertEqual(len(taskDefs), 1)
277 self.assertEqual(taskDefs[0].config.addend, 1000)
279 # Check --instrument option, for now it only checks that it does not crash
280 actions = [
281 _ACTION_ADD_TASK("testUtil.AddTask:task"),
282 _ACTION_ADD_INSTRUMENT("Instrument")
283 ]
284 args = _makeArgs(pipeline_actions=actions)
285 pipeline = fwk.makePipeline(args)
287 def testMakeGraphFromPickle(self):
288 """Tests for CmdLineFwk.makeGraph method.
290 Only most trivial case is tested that does not do actual graph
291 building.
292 """
293 fwk = CmdLineFwk()
295 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig:
297 # make non-empty graph and store it in a file
298 qgraph = _makeQGraph()
299 with open(tmpname, "wb") as pickleFile:
300 qgraph.save(pickleFile)
301 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
302 qgraph = fwk.makeGraph(None, args)
303 self.assertIsInstance(qgraph, QuantumGraph)
304 self.assertEqual(len(qgraph), 1)
306 # pickle with wrong object type
307 with open(tmpname, "wb") as pickleFile:
308 pickle.dump({}, pickleFile)
309 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
310 with self.assertRaises(TypeError):
311 fwk.makeGraph(None, args)
313 # reading empty graph from pickle should work but makeGraph()
314 # will return None and make a warning
315 qgraph = QuantumGraph()
316 with open(tmpname, "wb") as pickleFile:
317 qgraph.save(pickleFile)
318 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
319 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
320 # this also tests that warning is generated for empty graph
321 qgraph = fwk.makeGraph(None, args)
322 self.assertIs(qgraph, None)
324 def testSimpleQGraph(self):
325 """Test successfull execution of trivial quantum graph.
326 """
328 nQuanta = 5
329 butler, qgraph = makeSimpleQGraph(nQuanta)
331 # should have one task and number of quanta
332 self.assertEqual(len(qgraph), 1)
333 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
335 args = _makeArgs()
336 fwk = CmdLineFwk()
337 taskFactory = AddTaskFactoryMock()
339 # run whole thing
340 AddTask.countExec = 0
341 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
342 self.assertEqual(AddTask.countExec, nQuanta)
344 def testSimpleQGraphSkipExisting(self):
345 """Test continuing execution of trivial quantum graph with --skip-existing.
346 """
348 nQuanta = 5
349 butler, qgraph = makeSimpleQGraph(nQuanta)
351 # should have one task and number of quanta
352 self.assertEqual(len(qgraph), 1)
353 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
355 args = _makeArgs()
356 fwk = CmdLineFwk()
357 taskFactory = AddTaskFactoryMock()
359 # run whole thing
360 AddTask.countExec = 0
361 AddTask.stopAt = 3
362 with self.assertRaises(RuntimeError):
363 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
364 self.assertEqual(AddTask.countExec, 3)
366 AddTask.stopAt = -1
367 args.skip_existing = True
368 args.no_versions = True
369 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
370 self.assertEqual(AddTask.countExec, nQuanta)
372 def testShowPipeline(self):
373 """Test for --show options for pipeline.
374 """
375 fwk = CmdLineFwk()
377 actions = [
378 _ACTION_ADD_TASK("testUtil.AddTask:task"),
379 _ACTION_CONFIG("task:addend=100")
380 ]
381 args = _makeArgs(pipeline_actions=actions)
382 pipeline = fwk.makePipeline(args)
384 args.show = ["pipeline"]
385 fwk.showInfo(args, pipeline)
386 args.show = ["config"]
387 fwk.showInfo(args, pipeline)
388 args.show = ["history=task::addend"]
389 fwk.showInfo(args, pipeline)
390 args.show = ["tasks"]
391 fwk.showInfo(args, pipeline)
393 def testShowGraph(self):
394 """Test for --show options for quantum graph.
395 """
396 fwk = CmdLineFwk()
398 nQuanta = 2
399 butler, qgraph = makeSimpleQGraph(nQuanta)
401 args = _makeArgs()
402 args.show = ["graph"]
403 fwk.showInfo(args, pipeline=None, graph=qgraph)
404 # TODO: cannot test "workflow" option presently, it instanciates
405 # butler from command line options and there is no way to pass butler
406 # mock to that code.
409class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
410 pass
413def setup_module(module):
414 lsst.utils.tests.init()
417if __name__ == "__main__": 417 ↛ 418line 417 didn't jump to line 418, because the condition on line 417 was never true
418 lsst.utils.tests.init()
419 unittest.main()