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 args.timeout = None
188 return args
191def _makeQGraph():
192 """Make a trivial QuantumGraph with one quantum.
194 The only thing that we need to do with this quantum graph is to pickle
195 it, the quanta in this graph are not usable for anything else.
197 Returns
198 -------
199 qgraph : `~lsst.pipe.base.QuantumGraph`
200 """
201 taskDef = TaskDef(taskName="taskOne", config=SimpleConfig())
202 quanta = [Quantum()]
203 taskNodes = QuantumGraphTaskNodes(taskDef=taskDef, quanta=quanta, initInputs={}, initOutputs={})
204 qgraph = QuantumGraph([taskNodes])
205 return qgraph
208class CmdLineFwkTestCase(unittest.TestCase):
209 """A test case for CmdLineFwk
210 """
212 def testMakePipeline(self):
213 """Tests for CmdLineFwk.makePipeline method
214 """
215 fwk = CmdLineFwk()
217 # make empty pipeline
218 args = _makeArgs()
219 pipeline = fwk.makePipeline(args)
220 self.assertIsInstance(pipeline, Pipeline)
221 self.assertEqual(len(pipeline), 0)
223 # few tests with serialization
224 with makeTmpFile() as tmpname:
225 # make empty pipeline and store it in a file
226 args = _makeArgs(save_pipeline=tmpname)
227 pipeline = fwk.makePipeline(args)
228 self.assertIsInstance(pipeline, Pipeline)
230 # read pipeline from a file
231 args = _makeArgs(pipeline=tmpname)
232 pipeline = fwk.makePipeline(args)
233 self.assertIsInstance(pipeline, Pipeline)
234 self.assertEqual(len(pipeline), 0)
236 # single task pipeline
237 actions = [
238 _ACTION_ADD_TASK("TaskOne:task1")
239 ]
240 args = _makeArgs(pipeline_actions=actions)
241 pipeline = fwk.makePipeline(args)
242 self.assertIsInstance(pipeline, Pipeline)
243 self.assertEqual(len(pipeline), 1)
245 # many task pipeline
246 actions = [
247 _ACTION_ADD_TASK("TaskOne:task1a"),
248 _ACTION_ADD_TASK("TaskTwo:task2"),
249 _ACTION_ADD_TASK("TaskOne:task1b")
250 ]
251 args = _makeArgs(pipeline_actions=actions)
252 pipeline = fwk.makePipeline(args)
253 self.assertIsInstance(pipeline, Pipeline)
254 self.assertEqual(len(pipeline), 3)
256 # single task pipeline with config overrides, cannot use TaskOne, need
257 # something that can be imported with `doImport()`
258 actions = [
259 _ACTION_ADD_TASK("testUtil.AddTask:task"),
260 _ACTION_CONFIG("task:addend=100")
261 ]
262 args = _makeArgs(pipeline_actions=actions)
263 pipeline = fwk.makePipeline(args)
264 taskDefs = list(pipeline.toExpandedPipeline())
265 self.assertEqual(len(taskDefs), 1)
266 self.assertEqual(taskDefs[0].config.addend, 100)
268 overrides = b"config.addend = 1000\n"
269 with makeTmpFile(overrides) as tmpname:
270 actions = [
271 _ACTION_ADD_TASK("testUtil.AddTask:task"),
272 _ACTION_CONFIG_FILE("task:" + tmpname)
273 ]
274 args = _makeArgs(pipeline_actions=actions)
275 pipeline = fwk.makePipeline(args)
276 taskDefs = list(pipeline.toExpandedPipeline())
277 self.assertEqual(len(taskDefs), 1)
278 self.assertEqual(taskDefs[0].config.addend, 1000)
280 # Check --instrument option, for now it only checks that it does not crash
281 actions = [
282 _ACTION_ADD_TASK("testUtil.AddTask:task"),
283 _ACTION_ADD_INSTRUMENT("Instrument")
284 ]
285 args = _makeArgs(pipeline_actions=actions)
286 pipeline = fwk.makePipeline(args)
288 def testMakeGraphFromPickle(self):
289 """Tests for CmdLineFwk.makeGraph method.
291 Only most trivial case is tested that does not do actual graph
292 building.
293 """
294 fwk = CmdLineFwk()
296 with makeTmpFile() as tmpname, makeSQLiteRegistry() as registryConfig:
298 # make non-empty graph and store it in a file
299 qgraph = _makeQGraph()
300 with open(tmpname, "wb") as pickleFile:
301 qgraph.save(pickleFile)
302 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
303 qgraph = fwk.makeGraph(None, args)
304 self.assertIsInstance(qgraph, QuantumGraph)
305 self.assertEqual(len(qgraph), 1)
307 # pickle with wrong object type
308 with open(tmpname, "wb") as pickleFile:
309 pickle.dump({}, pickleFile)
310 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
311 with self.assertRaises(TypeError):
312 fwk.makeGraph(None, args)
314 # reading empty graph from pickle should work but makeGraph()
315 # will return None and make a warning
316 qgraph = QuantumGraph()
317 with open(tmpname, "wb") as pickleFile:
318 qgraph.save(pickleFile)
319 args = _makeArgs(qgraph=tmpname, registryConfig=registryConfig)
320 with self.assertWarnsRegex(UserWarning, "QuantumGraph is empty"):
321 # this also tests that warning is generated for empty graph
322 qgraph = fwk.makeGraph(None, args)
323 self.assertIs(qgraph, None)
325 def testSimpleQGraph(self):
326 """Test successfull execution of trivial quantum graph.
327 """
329 nQuanta = 5
330 butler, qgraph = makeSimpleQGraph(nQuanta)
332 # should have one task and number of quanta
333 self.assertEqual(len(qgraph), 1)
334 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
336 args = _makeArgs()
337 fwk = CmdLineFwk()
338 taskFactory = AddTaskFactoryMock()
340 # run whole thing
341 AddTask.countExec = 0
342 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
343 self.assertEqual(AddTask.countExec, nQuanta)
345 def testSimpleQGraphSkipExisting(self):
346 """Test continuing execution of trivial quantum graph with --skip-existing.
347 """
349 nQuanta = 5
350 butler, qgraph = makeSimpleQGraph(nQuanta)
352 # should have one task and number of quanta
353 self.assertEqual(len(qgraph), 1)
354 self.assertEqual(len(list(qgraph.quanta())), nQuanta)
356 args = _makeArgs()
357 fwk = CmdLineFwk()
358 taskFactory = AddTaskFactoryMock()
360 # run whole thing
361 AddTask.countExec = 0
362 AddTask.stopAt = 3
363 with self.assertRaises(RuntimeError):
364 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
365 self.assertEqual(AddTask.countExec, 3)
367 AddTask.stopAt = -1
368 args.skip_existing = True
369 args.no_versions = True
370 fwk.runPipeline(qgraph, taskFactory, args, butler=butler)
371 self.assertEqual(AddTask.countExec, nQuanta)
373 def testShowPipeline(self):
374 """Test for --show options for pipeline.
375 """
376 fwk = CmdLineFwk()
378 actions = [
379 _ACTION_ADD_TASK("testUtil.AddTask:task"),
380 _ACTION_CONFIG("task:addend=100")
381 ]
382 args = _makeArgs(pipeline_actions=actions)
383 pipeline = fwk.makePipeline(args)
385 args.show = ["pipeline"]
386 fwk.showInfo(args, pipeline)
387 args.show = ["config"]
388 fwk.showInfo(args, pipeline)
389 args.show = ["history=task::addend"]
390 fwk.showInfo(args, pipeline)
391 args.show = ["tasks"]
392 fwk.showInfo(args, pipeline)
394 def testShowGraph(self):
395 """Test for --show options for quantum graph.
396 """
397 fwk = CmdLineFwk()
399 nQuanta = 2
400 butler, qgraph = makeSimpleQGraph(nQuanta)
402 args = _makeArgs()
403 args.show = ["graph"]
404 fwk.showInfo(args, pipeline=None, graph=qgraph)
405 # TODO: cannot test "workflow" option presently, it instanciates
406 # butler from command line options and there is no way to pass butler
407 # mock to that code.
410class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase):
411 pass
414def setup_module(module):
415 lsst.utils.tests.init()
418if __name__ == "__main__": 418 ↛ 419line 418 didn't jump to line 419, because the condition on line 418 was never true
419 lsst.utils.tests.init()
420 unittest.main()