Coverage for tests/test_repository.py : 16%

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# -*- coding: UTF-8 -*-
3#
4# LSST Data Management System
5# Copyright 2016 LSST Corporation.
6#
7# This product includes software developed by the
8# LSST Project (http://www.lsst.org/).
9#
10# This program is free software: you can redistribute it and/or modify
11# it under the terms of the GNU General Public License as published by
12# the Free Software Foundation, either version 3 of the License, or
13# (at your option) any later version.
14#
15# This program is distributed in the hope that it will be useful,
16# but WITHOUT ANY WARRANTY; without even the implied warranty of
17# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
18# GNU General Public License for more details.
19#
20# You should have received a copy of the LSST License Statement and
21# the GNU General Public License along with this program. If not,
22# see <http://www.lsstcorp.org/LegalNotices/>.
23#
25import os
26import astropy.io.fits
27import shutil
28import sqlite3
29import unittest
30import tempfile
32import lsst.daf.persistence as dp
33# can't use name TestObject, becuase it messes Pytest up. Alias it to tstObj
34from lsst.daf.persistence.test import TestObject as tstObj
35from lsst.daf.persistence.test import MapperForTestWriting
36import lsst.utils.tests
38# Define the root of the tests relative to this file
39ROOT = os.path.abspath(os.path.dirname(__file__))
42def setup_module(module):
43 lsst.utils.tests.init()
46class ParentMapper(dp.Mapper):
48 def __init__(self, root, **kwargs):
49 self.root = root
50 self.storage = dp.Storage.makeFromURI(self.root)
52 def __repr__(self):
53 return 'ParentMapper(root=%s)' % self.root
55 def map_raw(self, dataId, write):
56 python = 'astropy.io.fits.HDUList'
57 persistable = None
58 storage = 'PickleStorage'
59 path = os.path.join(self.root, 'raw')
60 path = os.path.join(path, 'raw_v' + str(dataId['visit']) + '_f' + dataId['filter'] + '.fits.gz')
61 if os.path.exists(path):
62 return dp.ButlerLocation(python, persistable, storage, path,
63 dataId, self, self.storage)
64 return None
66 def bypass_raw(self, datasetType, pythonType, location, dataId):
67 return astropy.io.fits.open(location.getLocations()[0])
69 def query_raw(self, format, dataId):
70 values = [{'visit': 1, 'filter': 'g'}, {'visit': 2, 'filter': 'g'}, {'visit': 3, 'filter': 'r'}]
71 matches = []
72 for value in values:
73 match = True
74 for item in dataId:
75 if value[item] != dataId[item]:
76 match = False
77 break
78 if match:
79 matches.append(value)
80 results = set()
81 for match in matches:
82 tempTup = []
83 for word in format:
84 tempTup.append(match[word])
85 results.add(tuple(tempTup))
86 return results
88 def getDefaultLevel(self):
89 return 'visit'
91 def getKeys(self, datasetType, level):
92 return {'filter': str, 'visit': int}
94 def map_str(self, dataId, write):
95 path = os.path.join(self.root, 'raw')
96 path = os.path.join(path, 'raw_v' + str(dataId['str']) + '_f' + dataId['filter'] + '.fits.gz')
97 if os.path.exists(path):
98 return dp.ButlerLocation(str, None, 'PickleStorage', path, dataId,
99 self, self.storage)
100 return None
103class ChildrenMapper(dp.Mapper):
105 def __init__(self, root, **kwargs):
106 self.root = root
107 self.storage = dp.Storage.makeFromURI(self.root)
109 def map_raw(self, dataId, write):
110 python = 'astropy.io.fits.HDUList'
111 persistable = None
112 storage = 'FitsStorage'
113 path = os.path.join(self.root, 'raw')
114 path = os.path.join(path, 'raw_v' + str(dataId['visit']) + '_f' + dataId['filter'] + '.fits.gz')
115 if write or os.path.exists(path):
116 return dp.ButlerLocation(python, persistable, storage, path,
117 dataId, self, self.storage)
118 return None
120 def map_pickled(self, dataId, write):
121 python = 'dict'
122 persistable = None
123 storage = 'PickleStorage'
124 path = os.path.join(self.root, 'raw')
125 path = os.path.join(path, 'pickled_v' + str(dataId['visit']) + '_f' + dataId['filter'] + '.fits.gz')
126 if write or os.path.exists(path):
127 return dp.ButlerLocation(python, persistable, storage, path,
128 dataId, self, self.storage)
130 def bypass_raw(self, datasetType, pythonType, location, dataId):
131 return astropy.io.fits.open(location.getLocations()[0])
133 def query_raw(self, format, dataId):
134 return None
135 # results = set()
136 # return results
138 def getDefaultLevel(self):
139 return 'visit'
141 def getKeys(self, datasetType, level):
142 return {'filter': str, 'visit': int}
145class TestBasics(unittest.TestCase):
146 """Test case for basic functions of the repository classes."""
148 def setUp(self):
149 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestBasics-")
150 inputRepoArgs = dp.RepositoryArgs(root=os.path.join(ROOT, 'butlerAlias', 'data', 'input'),
151 mapper=ParentMapper,
152 tags='baArgs')
153 # mode of output repos is write-only by default
154 outputRepoArgs = dp.RepositoryArgs(root=os.path.join(self.testDir, 'repoA'),
155 mapper=ChildrenMapper,
156 mode='rw')
157 self.butler = dp.Butler(inputs=inputRepoArgs, outputs=outputRepoArgs)
159 def tearDown(self):
160 if os.path.exists(self.testDir):
161 shutil.rmtree(self.testDir)
162 if os.path.exists(os.path.join(ROOT, 'butlerAlias/repositoryCfg.yaml')):
163 os.remove(os.path.join(ROOT, 'butlerAlias/repositoryCfg.yaml'))
164 if hasattr(self, "butler"):
165 del self.butler
167 def testGet(self):
168 raw_image = self.butler.get('raw', {'visit': '2', 'filter': 'g'})
169 # in this case the width is known to be 1026:
170 self.assertEqual(raw_image[1].header["NAXIS1"], 1026)
172 def testGetUri(self):
173 raw_uri = self.butler.getUri('raw', {'visit': '2', 'filter': 'g'})
174 self.assertEqual(raw_uri,
175 os.path.join(ROOT, 'butlerAlias', 'data', 'input', 'raw', 'raw_v2_fg.fits.gz'))
176 self.assertEqual(os.path.isfile(raw_uri), True)
178 def testGetUriWrite(self):
179 raw_uri = self.butler.getUri('raw', {'visit': '9', 'filter': 'g'}, write=True)
180 self.assertEqual(raw_uri,
181 os.path.join(self.testDir, 'repoA', 'raw', 'raw_v9_fg.fits.gz'))
182 self.assertEqual(os.path.isfile(raw_uri), False)
184 def testSubset(self):
185 subset = self.butler.subset('raw')
186 self.assertEqual(len(subset), 3)
188 def testGetKeys(self):
189 keys = self.butler.getKeys('raw')
190 self.assertIn('filter', keys)
191 self.assertIn('visit', keys)
192 self.assertEqual(keys['filter'], str)
193 self.assertEqual(keys['visit'], int)
195 keys = self.butler.getKeys('raw', tag='baArgs')
196 self.assertIn('filter', keys)
197 self.assertIn('visit', keys)
198 self.assertEqual(keys['filter'], str)
199 self.assertEqual(keys['visit'], int)
201 keys = self.butler.getKeys('raw', tag=('baArgs', 'foo'))
202 self.assertIn('filter', keys)
203 self.assertIn('visit', keys)
204 self.assertEqual(keys['filter'], str)
205 self.assertEqual(keys['visit'], int)
207 keys = self.butler.getKeys('raw', tag='foo')
208 self.assertIsNone(keys)
210 keys = self.butler.getKeys('raw', tag=set(['baArgs', 'foo']))
211 self.assertIn('filter', keys)
212 self.assertIn('visit', keys)
213 self.assertEqual(keys['filter'], str)
214 self.assertEqual(keys['visit'], int)
216 def testQueryMetadata(self):
217 val = self.butler.queryMetadata('raw', ('filter',))
218 val.sort()
219 self.assertEqual(val, ['g', 'r'])
221 val = self.butler.queryMetadata('raw', ('visit',))
222 val.sort()
223 self.assertEqual(val, [1, 2, 3])
225 val = self.butler.queryMetadata('raw', ('visit',), dataId={'filter': 'g'})
226 val.sort()
227 self.assertEqual(val, [1, 2])
229 val = self.butler.queryMetadata('raw', ('visit',), dataId={'filter': 'r'})
230 self.assertEqual(val, [3])
232 val = self.butler.queryMetadata('raw', ('filter',), dataId={'visit': 1})
233 self.assertEqual(val, ['g'])
235 val = self.butler.queryMetadata('raw', ('filter',), dataId={'visit': 2})
236 self.assertEqual(val, ['g'])
238 val = self.butler.queryMetadata('raw', ('filter',), dataId={'visit': 3})
239 self.assertEqual(val, ['r'])
241 val = self.butler.queryMetadata('raw', ('visit',), dataId={'filter': 'h'})
242 self.assertEqual(val, [])
244 dataId = dp.DataId({'filter': 'g'}, tag='baArgs')
245 val = self.butler.queryMetadata('raw', ('visit',), dataId={'filter': 'g'})
246 val.sort()
247 self.assertEqual(val, [1, 2])
249 dataId = dp.DataId({'filter': 'g'}, tag='foo')
250 val = self.butler.queryMetadata('raw', ('visit',), dataId=dataId)
251 self.assertEqual(val, [])
253 def testDatasetExists(self):
254 # test the values that are expected to be true:
255 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'g', 'visit': 1}), True)
256 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'g', 'visit': 2}), True)
257 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'r', 'visit': 3}), True)
259 # test a few values that are expected to be false:
260 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'f', 'visit': 1}), False)
261 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'r', 'visit': 1}), False)
262 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'g', 'visit': 3}), False)
264 # test that we don't see datasets in the input repo as being in the output repo
265 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'g', 'visit': 1}, write=True), False)
266 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'g', 'visit': 2}, write=True), False)
267 self.assertEqual(self.butler.datasetExists('raw', {'filter': 'r', 'visit': 3}, write=True), False)
269 # write a dataset to the output repo and check that we see it with both write=True and write=False
270 dataId = {'visit': '2', 'filter': 'g'}
271 self.butler.put({"number": 3}, 'pickled', dataId)
272 self.assertEqual(self.butler.datasetExists('pickled', dataId, write=True), True)
273 self.assertEqual(self.butler.datasetExists('pickled', dataId, write=False), True)
276##############################################################################################################
277##############################################################################################################
278##############################################################################################################
281class NoResultsMapper(dp.Mapper):
283 def __init__(self, root, **kwargs):
284 self.root = root
285 self.storage = dp.Storage.makeFromURI(self.root)
287 def __repr__(self):
288 return 'ParentMapper(root=%s)' % self.root
290 def map_list(self, dataId, write):
291 return []
293 def map_None(self, dataId, write):
294 return None
297class TestWriting(unittest.TestCase):
298 """A test case for the repository classes.
300 A test that
301 1. creates repo with a peer repo, writes to those repos.
302 2. reloads those output repos as a parents of new repos
303 * does a read from from the repo (verifies parent search)
304 3. writes to the new output repo and reloads it as a parent of a new repo
305 * verifies masking
306 4. reloads the repo from its persisted cfg
307 * verifies reload from cfg
308 """
309 def setUp(self):
310 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestWriting-")
312 def tearDown(self):
313 if os.path.exists(self.testDir):
314 shutil.rmtree(self.testDir)
315 # del self.butler
317 def testCreateAggregateAndLoadingAChild(self):
318 """Tests putting a very basic pickled object in a variety of Repository configuration settings
319 :return:
320 """
322 repoAArgs = dp.RepositoryArgs(mode='w',
323 root=os.path.join(self.testDir, 'repoA'),
324 mapper=MapperForTestWriting)
325 repoBArgs = dp.RepositoryArgs(mode='w',
326 root=os.path.join(self.testDir, 'repoB'),
327 mapper=MapperForTestWriting)
328 butlerAB = dp.Butler(outputs=[repoAArgs, repoBArgs])
330 objA = tstObj('abc')
331 butlerAB.put(objA, 'foo', {'val': 1})
332 objB = tstObj('def')
333 butlerAB.put(objB, 'foo', {'val': 2})
335 # create butlers where the output repos are now input repos
336 butlerC = dp.Butler(inputs=dp.RepositoryArgs(root=os.path.join(self.testDir, 'repoA')))
337 butlerD = dp.Butler(inputs=dp.RepositoryArgs(root=os.path.join(self.testDir, 'repoB')))
339 # # verify the objects exist by getting them
340 self.assertEqual(objA, butlerC.get('foo', {'val': 1}))
341 self.assertEqual(objA, butlerC.get('foo', {'val': 1}))
342 self.assertEqual(objB, butlerD.get('foo', {'val': 2}))
343 self.assertEqual(objB, butlerD.get('foo', {'val': 2}))
345 def testPutWithIncompleteDataId(self):
346 repoAArgs = dp.RepositoryArgs(mode='rw',
347 root=os.path.join(self.testDir, 'repoA'),
348 mapper=NoResultsMapper)
349 butler = dp.Butler(outputs=repoAArgs)
350 obj = tstObj('abc')
351 with self.assertRaises(RuntimeError):
352 butler.put(obj, 'list', {'visit': '2'})
353 with self.assertRaises(RuntimeError):
354 butler.put(obj, 'None', {'visit': '2'})
357class TestDiamondPattern(unittest.TestCase):
358 """A test case for when a butler is created with an output repository and two input repositories that have
359 a common parent; verifies that the parent's registry is loaded as the parentRegistry of the output
360 repository.
361 """
362 @staticmethod
363 def makeRegistry(location, name):
364 conn = sqlite3.connect(location)
365 conn.execute("CREATE TABLE {name} (real)".format(name=name))
366 conn.commit()
367 conn.close()
369 def setUp(self):
370 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestDiamondPattern-")
372 def tearDown(self):
373 if os.path.exists(self.testDir):
374 shutil.rmtree(self.testDir)
376 def test(self):
377 repoARoot = os.path.join(self.testDir, 'repoA')
378 butler = dp.Butler(outputs={'root': repoARoot, 'mapper': ParentRepoTestMapper})
379 self.makeRegistry(os.path.join(repoARoot, 'registry.sqlite3'), 'repoA')
380 del butler
382 repoBRoot = os.path.join(self.testDir, 'repoB')
383 butlerB = dp.Butler(inputs=repoARoot, outputs=repoBRoot)
384 del butlerB
386 repoCRoot = os.path.join(self.testDir, 'repoC')
387 butlerC = dp.Butler(inputs=repoARoot, outputs=repoCRoot)
388 del butlerC
390 repoDRoot = os.path.join(self.testDir, 'repoD')
391 butlerD = dp.Butler(inputs=(repoBRoot, repoCRoot), outputs=repoDRoot)
393 self.assertEqual(1, len(butlerD._repos.outputs()))
394 self.assertEqual(os.path.dirname(butlerD._repos.outputs()[0].parentRegistry.root), repoARoot)
397class TestMasking(unittest.TestCase):
398 """A test case for the repository classes.
400 A test that
401 1. creates a repo, does a put
402 2. creates a new butler, uses that repo as an input and creates a new read-write output repo
403 3. gets from the input repo, modifies the dataset, and puts into the output repo
404 4. does a get and verifies that the changed dataset is returned.
405 """
406 def setUp(self):
407 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestMasking-")
409 def tearDown(self):
410 if os.path.exists(self.testDir):
411 shutil.rmtree(self.testDir)
413 def test(self):
414 repoAArgs = dp.RepositoryArgs(mode='w',
415 root=os.path.join(self.testDir, 'repoA'),
416 mapper=MapperForTestWriting)
417 butler = dp.Butler(outputs=repoAArgs)
418 obj0 = tstObj('abc')
419 butler.put(obj0, 'foo', {'bar': 1})
420 del butler
421 del repoAArgs
423 repoBArgs = dp.RepositoryArgs(mode='rw',
424 root=os.path.join(self.testDir, 'repoB'),
425 mapper=MapperForTestWriting)
426 butler = dp.Butler(inputs=os.path.join(self.testDir, 'repoA'), outputs=repoBArgs)
427 obj1 = butler.get('foo', {'bar': 1})
428 self.assertEqual(obj0, obj1)
429 obj1.data = "def"
430 butler.put(obj1, 'foo', {'bar': 1})
431 obj2 = butler.get('foo', {'bar': 1})
432 self.assertEqual(obj1, obj2)
435class TestMultipleOutputsPut(unittest.TestCase):
436 """A test case for the repository classes.
438 A test that
439 1. creates 3 peer repositories and readers for them
440 2. does a single put
441 3. verifies that all repos received the put
442 """
443 def setUp(self):
444 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestMultipleOutputsPut-")
446 def tearDown(self):
447 if os.path.exists(self.testDir):
448 shutil.rmtree(self.testDir)
450 def test(self):
451 repoAArgs = dp.RepositoryArgs(root=os.path.join(self.testDir, 'repoA'),
452 mapper=MapperForTestWriting)
453 repoBArgs = dp.RepositoryArgs(root=os.path.join(self.testDir, 'repoB'),
454 mapper=MapperForTestWriting)
456 butler = dp.Butler(outputs=(repoAArgs, repoBArgs))
457 obj0 = tstObj('abc')
458 butler.put(obj0, 'foo', {'bar': 1})
460 for root in (os.path.join(self.testDir, 'repoA'),
461 os.path.join(self.testDir, 'repoB')):
462 butler = dp.Butler(inputs=root)
463 self.assertEqual(butler.get('foo', {'bar': 1}), obj0)
466class TestMultipleInputs(unittest.TestCase):
467 """A test case for the repository classes.
469 A test that
470 - create output 2 repos, write same & different data to them & close them
471 - create a new butler with those repos as inputs
472 - read data that was written to both repos:
473 - verify data that was written to only one of each repo
474 - verify dissimilar datasets with same id that were written to the repos
476 """
477 def setUp(self):
478 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestMapperInference-")
480 def tearDown(self):
481 if os.path.exists(self.testDir):
482 shutil.rmtree(self.testDir)
484 def test(self):
485 objAbc = tstObj('abc')
486 objDef = tstObj('def')
487 objGhi = tstObj('ghi')
488 objJkl = tstObj('jkl')
490 repoAArgs = dp.RepositoryArgs(mode='w',
491 root=os.path.join(self.testDir, 'repoA'),
492 mapper=MapperForTestWriting)
493 butler = dp.Butler(outputs=repoAArgs)
494 butler.put(objAbc, 'foo', {'bar': 1})
495 butler.put(objDef, 'foo', {'bar': 2})
497 repoBArgs = dp.RepositoryArgs(mode='w',
498 root=os.path.join(self.testDir, 'repoB'),
499 mapper=MapperForTestWriting)
500 butler = dp.Butler(outputs=repoBArgs)
501 butler.put(objGhi, 'foo', {'bar': 1}) # note different object with overlapping dataId with repoA
502 butler.put(objJkl, 'foo', {'bar': 3})
504 # note repo order: A, B
505 butler = dp.Butler(inputs=(os.path.join(self.testDir, 'repoA'),
506 os.path.join(self.testDir, 'repoB')))
507 self.assertEqual(butler.get('foo', {'bar': 1}), objAbc)
508 self.assertEqual(butler.get('foo', {'bar': 2}), objDef)
509 self.assertEqual(butler.get('foo', {'bar': 3}), objJkl)
511 # note reversed repo order: B, A
512 butler = dp.Butler(inputs=(os.path.join(self.testDir, 'repoB'),
513 os.path.join(self.testDir, 'repoA')))
514 self.assertEqual(butler.get('foo', {'bar': 1}), objGhi)
515 self.assertEqual(butler.get('foo', {'bar': 2}), objDef)
516 self.assertEqual(butler.get('foo', {'bar': 3}), objJkl)
519class TestTagging(unittest.TestCase):
520 """A test case for the tagging of repository classes.
521 """
523 def setUp(self):
524 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestTagging-")
526 def tearDown(self):
527 if os.path.exists(self.testDir):
528 shutil.rmtree(self.testDir)
530 def testOneLevelInputs(self):
531 """
532 1. put an object with the same ID but slightly different value into 2 repositories.
533 2. use those repositories as inputs to a butler, and tag them
534 3. make sure that the correct object is gotten for each of
535 a. one tag
536 b. the other tag
537 c. no tag
538 4. repeat step 3 but reverse the order of input cfgs to a new butler.
539 5. use the butler from step 4 and write an output. The inputs will get recorded as parents of the
540 output repo.
541 6. create a new butler with a new overlapping repo, and verify that objects can be gotten from the
542 other's parent repos via tagging.
543 """
544 objA = tstObj('a')
545 objB = tstObj('b')
547 # put objA in repo1:
548 repo1Args = dp.RepositoryArgs(mode='rw',
549 root=os.path.join(self.testDir, 'repo1'),
550 mapper=MapperForTestWriting)
551 butler = dp.Butler(outputs=repo1Args)
552 butler.put(objA, 'foo', {'bar': 1})
553 del butler
555 # put objB in repo2:
556 repo2Args = dp.RepositoryArgs(mode='rw',
557 root=os.path.join(self.testDir, 'repo2'),
558 mapper=MapperForTestWriting)
559 butler = dp.Butler(outputs=repo2Args)
560 butler.put(objB, 'foo', {'bar': 1})
561 del butler
562 del repo1Args
563 del repo2Args
565 # make the objects inputs of repos
566 # and verify the correct object can ge fetched using the tag and not using the tag
568 repo1Args = dp.RepositoryArgs(root=os.path.join(self.testDir, 'repo1'), tags='one')
569 repo2Args = dp.RepositoryArgs(root=os.path.join(self.testDir, 'repo2'), tags='two')
571 butler = dp.Butler(inputs=(repo1Args, repo2Args))
572 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='one')), objA)
573 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='two')), objB)
574 self.assertEqual(butler.get('foo', {'bar': 1}), objA)
576 butler = dp.Butler(inputs=(repo2Args, repo1Args))
577 self.assertEqual(butler.get('foo', dp.DataId(bar=1, tag='one')), objA)
578 self.assertEqual(butler.get('foo', dp.DataId(bar=1, tag='two')), objB)
579 self.assertEqual(butler.get('foo', dp.DataId(bar=1)), objB)
581 # create butler with repo1 and repo2 as parents, and an output repo3.
582 repo3Args = dp.RepositoryArgs(mode='rw',
583 root=os.path.join(self.testDir, 'repo3'),
584 mapper=MapperForTestWriting)
585 butler = dp.Butler(inputs=(repo1Args, repo2Args), outputs=repo3Args)
586 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='one')), objA)
587 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='two')), objB)
588 self.assertEqual(butler.get('foo', {'bar': 1}), objA)
589 # add an object to the output repo. note since the output repo mode is 'rw' that object is gettable
590 # and it has first priority in search order. Other repos should be searchable by tagging.
591 objC = tstObj('c')
592 butler.put(objC, 'foo', {'bar': 1})
593 self.assertEqual(butler.get('foo', {'bar': 1}), objC)
594 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='one')), objA)
595 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='two')), objB)
596 del butler
598 repo3Cfg = dp.Storage().getRepositoryCfg(os.path.join(self.testDir, 'repo3'))
599 self.assertEqual(repo3Cfg.parents, [os.path.join(self.testDir, 'repo1'),
600 os.path.join(self.testDir, 'repo2')])
602 # expand the structure to look like this:
603 # ┌────────────────────────┐ ┌────────────────────────┐
604 # │repo1 │ │repo2 │
605 # │ tag:"one" │ │ tag:"two" │
606 # │ tstObj('a') │ │ tstObj('b') │
607 # │ at ('foo', {'bar:1'})│ │ at ('foo', {'bar:1'})│
608 # └───────────┬────────────┘ └───────────┬────────────┘
609 # └─────────────┬────────────┘
610 # ┌────────────┴───────────┐ ┌────────────────────────┐
611 # │repo4 │ │repo5 │
612 # │ tag:"four" │ │ tag:"five" │
613 # │ tstObj('d') │ │ tstObj('e') │
614 # │ at ('foo', {'bar:2'})│ │ at ('foo', {'bar:1'})│
615 # └───────────┬────────────┘ └───────────┬────────────┘
616 # └─────────────┬────────────┘
617 # ┌──┴───┐
618 # │butler│
619 # └──────┘
621 repo4Args = dp.RepositoryArgs(mode='rw',
622 root=os.path.join(self.testDir, 'repo4'),
623 mapper=MapperForTestWriting)
624 butler = dp.Butler(inputs=(os.path.join(self.testDir, 'repo1'),
625 os.path.join(self.testDir, 'repo2')), outputs=repo4Args)
626 objD = tstObj('d')
627 butler.put(objD, 'foo', {'bar': 2})
628 del butler
630 repo5Cfg = dp.RepositoryArgs(mode='rw',
631 root=os.path.join(self.testDir, 'repo5'),
632 mapper=MapperForTestWriting)
633 butler = dp.Butler(outputs=repo5Cfg)
634 objE = tstObj('e')
635 butler.put(objE, 'foo', {'bar': 1})
636 del butler
638 repo4Args = dp.RepositoryArgs(cfgRoot=os.path.join(self.testDir, 'repo4'), tags='four')
639 repo5Args = dp.RepositoryArgs(cfgRoot=os.path.join(self.testDir, 'repo5'), tags='five')
640 butler = dp.Butler(inputs=(repo4Args, repo5Args))
641 self.assertEqual(butler.get('foo', {'bar': 1}), objA)
642 self.assertEqual(butler.get('foo', {'bar': 2}), objD)
643 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='four')), objA)
644 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='five')), objE)
645 del butler
647 butler = dp.Butler(inputs=(repo5Args, repo4Args))
648 self.assertEqual(butler.get('foo', {'bar': 1}), objE)
649 self.assertEqual(butler.get('foo', {'bar': 2}), objD)
650 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='four')), objA)
651 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='five')), objE)
652 del butler
655class TestMapperInference(unittest.TestCase):
656 """A test for inferring mapper in the cfg from parent cfgs"""
658 def setUp(self):
659 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestMapperInference-")
661 def tearDown(self):
662 if os.path.exists(self.testDir):
663 shutil.rmtree(self.testDir)
665 def testSingleParent(self):
666 """
667 creates a repo that:
668 a. does not have a mapper specified in the cfg
669 b. has a parent that does have a mapper specified in the cfg
670 verifies that the child repo inherits the parent's mapper via inference.
671 """
672 repoACfg = dp.RepositoryCfg(root=os.path.join(self.testDir, 'repoA'),
673 mapper=MapperForTestWriting,
674 mapperArgs=None,
675 parents=None,
676 policy=None)
677 dp.Storage.putRepositoryCfg(repoACfg, os.path.join(self.testDir, 'repoA'))
679 repoBArgs = dp.RepositoryArgs(mode='rw',
680 root=os.path.join(self.testDir, 'repoB'))
681 butler = dp.Butler(inputs=os.path.join(self.testDir, 'repoA'), outputs=repoBArgs)
682 self.assertIsInstance(butler._repos.outputs()[0].repo._mapper, MapperForTestWriting)
684 def testMultipleParentsSameMapper(self):
685 """
686 creates a repo that:
687 a. does not have a mapper specified in the cfg
688 b. has 2 parents that do have the same mapper specified in the cfg
689 verifies that the child repo inherits the parent's mapper via inference.
691 """
692 repoACfg = dp.RepositoryCfg(root=os.path.join(self.testDir, 'repoA'),
693 mapper=MapperForTestWriting,
694 mapperArgs=None,
695 parents=None,
696 policy=None,)
697 repoBCfg = dp.RepositoryCfg(root=os.path.join(self.testDir, 'repoB'),
698 mapper=MapperForTestWriting,
699 mapperArgs=None,
700 parents=None,
701 policy=None)
702 dp.Storage.putRepositoryCfg(repoACfg, os.path.join(self.testDir, 'repoA'))
703 dp.Storage.putRepositoryCfg(repoBCfg, os.path.join(self.testDir, 'repoB'))
705 repoCArgs = dp.RepositoryArgs(mode='w',
706 root=os.path.join(self.testDir, 'repoC'))
708 butler = dp.Butler(inputs=(os.path.join(self.testDir, 'repoA'),
709 os.path.join(self.testDir, 'repoB')),
710 outputs=repoCArgs)
711 self.assertIsInstance(butler._repos.outputs()[0].repo._mapper, MapperForTestWriting)
713 def testMultipleParentsDifferentMappers(self):
714 """
715 creates a repo that:
716 a. does not have a mapper specified in the cfg
717 b. has 2 parent repos that have different mappers
718 verifies that the constructor raises a RuntimeError because the mapper can not be inferred.
719 """
720 repoACfg = dp.RepositoryCfg(root=os.path.join(self.testDir, 'repoA'),
721 mapper=MapperForTestWriting,
722 mapperArgs=None,
723 parents=None,
724 policy=None)
725 repoBCfg = dp.RepositoryCfg(root=os.path.join(self.testDir, 'repoB'),
726 mapper='lsst.daf.persistence.test.EmptyTestMapper',
727 mapperArgs=None,
728 parents=None,
729 policy=None)
730 dp.Storage.putRepositoryCfg(repoACfg, os.path.join(self.testDir, 'repoA'))
731 dp.Storage.putRepositoryCfg(repoBCfg, os.path.join(self.testDir, 'repoB'))
733 repoCArgs = dp.RepositoryArgs(mode='w',
734 root=os.path.join(self.testDir, 'repoC'))
735 self.assertRaises(RuntimeError,
736 dp.Butler,
737 inputs=(os.path.join(self.testDir, 'repoA'),
738 os.path.join(self.testDir, 'repoB')),
739 outputs=repoCArgs)
742class TestMovedRepositoryCfg(unittest.TestCase):
743 """Test if a repository cfg is in-place (resides at root of the repository) and the cfg is moved, the
744 deserialized cfg root should be the new location if the repository is moved.
745 """
747 def setUp(self):
748 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestMovedRepositoryCfg-")
750 def tearDown(self):
751 if os.path.exists(self.testDir):
752 shutil.rmtree(self.testDir)
754 def test(self):
755 butler = dp.Butler(outputs=dp.RepositoryArgs(root=os.path.join(self.testDir, 'a'),
756 mapper=MapperForTestWriting))
757 del butler
758 os.makedirs(os.path.join(self.testDir, 'b'))
759 os.rename(os.path.join(self.testDir, 'a/repositoryCfg.yaml'),
760 os.path.join(self.testDir, 'b/repositoryCfg.yaml'))
761 butler = dp.Butler(inputs=os.path.join(self.testDir, 'b'))
762 self.assertEqual(list(butler._repos.inputs())[0].cfg,
763 dp.RepositoryCfg(root=os.path.join(self.testDir, 'b'),
764 mapper=MapperForTestWriting,
765 mapperArgs=None,
766 parents=None,
767 policy=None))
770class TestOutputAlreadyHasParent(unittest.TestCase):
772 def setUp(self):
773 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestOutputAlreadyHasParent-")
775 def tearDown(self):
776 if os.path.exists(self.testDir):
777 shutil.rmtree(self.testDir)
779 def test(self):
780 # create a repo where repo 'a' is a parent of repo 'b'
781 butler = dp.Butler(outputs=dp.RepositoryArgs(root=os.path.join(self.testDir, 'a'),
782 mapper=MapperForTestWriting))
783 del butler
784 butler = dp.Butler(inputs=os.path.join(self.testDir, 'a'),
785 outputs=os.path.join(self.testDir, 'b'))
786 self.assertEqual(len(butler._repos.inputs()), 1)
787 self.assertEqual(butler._repos.inputs()[0].cfg.root,
788 os.path.join(self.testDir, 'a'))
789 self.assertEqual(len(butler._repos.outputs()), 1)
790 self.assertEqual(butler._repos.outputs()[0].cfg.root,
791 os.path.join(self.testDir, 'b'))
792 del butler
794 # load that repo a few times, include 'a' as an input.
795 for i in range(4):
796 butler = dp.Butler(inputs=os.path.join(self.testDir, 'a'),
797 outputs=dp.RepositoryArgs(root=os.path.join(self.testDir, 'b'),
798 mode='rw'))
799 self.assertEqual(len(butler._repos.inputs()), 2)
800 self.assertEqual(butler._repos.inputs()[0].cfg.root,
801 os.path.join(self.testDir, 'b'))
802 self.assertEqual(butler._repos.inputs()[1].cfg.root,
803 os.path.join(self.testDir, 'a'))
804 self.assertEqual(len(butler._repos.outputs()), 1)
805 self.assertEqual(butler._repos.outputs()[0].cfg.root,
806 os.path.join(self.testDir, 'b'))
807 cfg = dp.Storage().getRepositoryCfg(os.path.join(self.testDir, 'b'))
808 self.assertEqual(cfg, dp.RepositoryCfg(root=os.path.join(self.testDir, 'b'),
809 mapper=MapperForTestWriting,
810 mapperArgs=None,
811 parents=[os.path.join(self.testDir, 'a')],
812 policy=None))
814 # load the repo a few times and don't explicitly list 'a' as an input
815 for i in range(4):
816 butler = dp.Butler(outputs=dp.RepositoryArgs(root=os.path.join(self.testDir, 'b'),
817 mode='rw'))
818 self.assertEqual(len(butler._repos.inputs()), 2)
819 self.assertEqual(butler._repos.inputs()[0].cfg.root,
820 os.path.join(self.testDir, 'b'))
821 self.assertEqual(butler._repos.inputs()[1].cfg.root,
822 os.path.join(self.testDir, 'a'))
823 self.assertEqual(len(butler._repos.outputs()), 1)
824 self.assertEqual(butler._repos.outputs()[0].cfg.root,
825 os.path.join(self.testDir, 'b'))
826 cfg = dp.Storage().getRepositoryCfg(os.path.join(self.testDir, 'b'))
827 self.assertEqual(cfg, dp.RepositoryCfg(root=os.path.join(self.testDir, 'b'),
828 mapper=MapperForTestWriting,
829 mapperArgs=None,
830 parents=[os.path.join(self.testDir, 'a')],
831 policy=None))
833 # load 'b' as 'write only' and don't list 'a' as an input. This should raise, because inputs must
834 # match readable outputs parents.
835 with self.assertRaises(RuntimeError):
836 butler = dp.Butler(outputs=os.path.join(self.testDir, 'b'))
838 # load 'b' as 'write only' and explicitly list 'a' as an input.
839 butler = dp.Butler(inputs=os.path.join(self.testDir, 'a'),
840 outputs=os.path.join(self.testDir, 'b'))
841 self.assertEqual(len(butler._repos.inputs()), 1)
842 self.assertEqual(len(butler._repos.outputs()), 1)
843 self.assertEqual(butler._repos.inputs()[0].cfg.root,
844 os.path.join(self.testDir, 'a'))
845 self.assertEqual(butler._repos.outputs()[0].cfg.root,
846 os.path.join(self.testDir, 'b'))
847 cfg = dp.Storage().getRepositoryCfg(os.path.join(self.testDir, 'b'))
850class ParentRepoTestMapper(dp.Mapper):
851 def __init__(self, parentRegistry, repositoryCfg, **kwargs):
852 self.parentRegistry = parentRegistry
853 root = repositoryCfg.root
854 if os.path.exists(os.path.join(root, 'registry.sqlite3')):
855 self.registry = dp.Registry.create(os.path.join(root, 'registry.sqlite3'))
856 else:
857 self.registry = None
859 def getRegistry(self):
860 return self.registry
863class TestParentRepository(unittest.TestCase):
864 """A test to verify that when a parent repository is used the correct one is found."""
866 def setUp(self):
867 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestParentRepository-")
869 def tearDown(self):
870 if os.path.exists(self.testDir):
871 shutil.rmtree(self.testDir, True)
873 def test(self):
875 """Tests 1. That an sqlite registry in a parent repo is used as the
876 registry in a child repo that does not have its own sqlite registry.
877 2. That when a repo has a parent and grandparent and only the
878 grandparent has an sqlite registry that the grandparent's registry is
879 used.
880 3. That when the parent and grandparent both have sqlite registries
881 that the parent's sqlite registry is used."""
882 # setup; create the parent repo
883 def makeRegistry(location, name):
884 conn = sqlite3.connect(location)
885 conn.execute("CREATE TABLE {name} (real)".format(name=name))
886 conn.commit()
887 conn.close()
889 repoADir = os.path.join(self.testDir, 'repoA')
890 repoBDir = os.path.join(self.testDir, 'repoB')
891 repoCDir = os.path.join(self.testDir, 'repoC')
892 butler = dp.Butler(outputs={'root': repoADir, 'mapper': ParentRepoTestMapper})
893 del butler
894 makeRegistry(os.path.join(repoADir, 'registry.sqlite3'), 'repoA')
895 # test 1:
896 butler = dp.Butler(inputs=repoADir, outputs=repoBDir)
897 registry = butler._repos.outputs()[0].repo._mapper.parentRegistry
898 self.assertIsInstance(registry, dp.SqliteRegistry)
899 tables = registry.conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
900 self.assertEqual(tables, [('repoA', )])
901 del butler
902 # test 2:
903 butler = dp.Butler(inputs=repoBDir, outputs=repoCDir)
904 registry = butler._repos.outputs()[0].repo._mapper.parentRegistry
905 self.assertIsInstance(registry, dp.SqliteRegistry)
906 tables = registry.conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
907 self.assertEqual(tables, [('repoA', )])
908 del butler
909 # test 3:
910 makeRegistry(os.path.join(repoBDir, 'registry.sqlite3'), 'repoB')
911 butler = dp.Butler(inputs=repoBDir, outputs=repoCDir)
912 registry = butler._repos.outputs()[0].repo._mapper.parentRegistry
913 self.assertIsInstance(registry, dp.SqliteRegistry)
914 tables = registry.conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()
915 self.assertEqual(tables, [('repoB', )])
918class TestOldButlerParent(unittest.TestCase):
919 """A test to verify that when a parent is an old butler repo that it still gets loaded correctly,
920 including mapperArgs."""
922 def setUp(self):
923 """Create temp test dir and create an OldButler repo at 'repoA'"""
924 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestOldButlerParent-")
925 self.repoADir = os.path.join(self.testDir, 'repoA')
926 os.makedirs(self.repoADir)
927 with open(os.path.join(self.repoADir, '_mapper'), 'w') as mapperFile:
928 mapperFile.write('lsst.daf.persistence.test.EmptyTestMapper')
930 def tearDown(self):
931 if os.path.exists(self.testDir):
932 shutil.rmtree(self.testDir, True)
934 def testOldButlerCfgRecordedInOutputs(self):
935 """Verify that when an Old Butler is used as an input with parameters such as mapperArgs, that those
936 parameters get recalled correctly (in the RepositoryCfg nested in the parents list of the output repo)
937 """
938 repoBDir = os.path.join(self.testDir, 'repoB')
939 # create a 'repoB' with the Old Butler 'repoA' as a parent, specify mapperArgs for 'repoA', and
940 # verify that the mapperArgs got passed to the repoA mapper.
941 butler = dp.Butler(inputs={'root': self.repoADir,
942 'mapperArgs': {'foo': 1}},
943 outputs={'root': repoBDir,
944 'mapperArgs': {'foo': 1},
945 'mode': 'rw'})
946 self.assertEqual(butler._repos.inputs()[0].repo._mapper.kwargs, {'foo': 1})
947 # init a Butler again, with exactly the same configuration as before.
948 butler = dp.Butler(inputs={'root': self.repoADir,
949 'mapperArgs': {'foo': 1}},
950 outputs={'root': repoBDir,
951 'mapperArgs': {'foo': 1},
952 'mode': 'rw'})
953 self.assertEqual(butler._repos.inputs()[0].repo._mapper.kwargs, {'foo': 1})
954 # try initializing the butler again with the 'repoB' output, but do not specify the mapperArgs in the
955 # repoA args. this should raise a runtime error because the input speficied this way does not match
956 # the cfg recorded in 'repoB'.
957 with self.assertRaises(RuntimeError):
958 butler = dp.Butler(inputs=self.repoADir,
959 outputs=repoBDir)
960 # initializing the butler, but only specifiying the root of 'repoB' (the rest of the cfg will be
961 # loaded from its `repositoryCfg` file), but fully specify the input as required.
962 butler = dp.Butler(inputs={'root': self.repoADir,
963 'mapperArgs': {'foo': 1}},
964 outputs=repoBDir)
965 self.assertEqual(butler._repos.inputs()[0].repo._mapper.kwargs, {'foo': 1})
966 # use 'repoB' as an input, and verify that 'repoA' is loaded, including the mapperArgs.
967 butler = dp.Butler(inputs=repoBDir)
968 self.assertEqual(butler._repos.inputs()[1].repo._mapper.kwargs, {'foo': 1})
971class TestOldButlerParentTagging(unittest.TestCase):
972 """A test to verify that when a parent is an old butler repo that its tagging gets loaded correctly"""
974 def setUp(self):
975 """Remove testDir from any previous runs and create an OldButler repo at repoA and an Old Butler repoB
976 with a _parent symlink to repoA.
977 """
978 self.testDir = tempfile.mkdtemp(dir=ROOT, prefix="TestOldButlerParentTagging-")
979 self.repoADir = os.path.join(self.testDir, 'repoA')
980 os.makedirs(self.repoADir)
981 with open(os.path.join(self.repoADir, '_mapper'), 'w') as mapperFile:
982 mapperFile.write('lsst.daf.persistence.test.MapperForTestWriting')
983 self.repoBDir = os.path.join(self.testDir, 'repoB')
984 os.makedirs(self.repoBDir)
985 os.symlink(self.repoADir, os.path.join(self.repoBDir, '_parent'))
987 def tearDown(self):
988 if os.path.exists(self.testDir):
989 shutil.rmtree(self.testDir, True)
991 def test(self):
992 """Verify that the tags on a repository with an Old Butler repository parent are applied to that
993 parent
994 """
995 # put objA in repoA:
996 objA = tstObj('a')
997 butler = dp.Butler(outputs=self.repoADir)
998 butler.put(objA, 'foo', {'bar': 1})
999 del butler
1001 # create repoB and put objB in it:
1002 objB = tstObj('b')
1003 butler = dp.Butler(inputs=self.repoADir, outputs=self.repoBDir)
1004 butler.put(objB, 'foo', {'bar': 2})
1005 del butler
1007 # verify that repoB can be used as an input, and objA and objB can be gotten from it:
1008 butler = dp.Butler(inputs=self.repoBDir)
1009 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1})), objA)
1010 self.assertEqual(butler.get('foo', dp.DataId({'bar': 2})), objB)
1011 del butler
1013 # apply a tag and verify that repoB can still be used as an input, and both objA (in repoA) and objB
1014 # can be gotten from it:
1015 butler = dp.Butler(inputs={'root': self.repoBDir, 'tags': 'baz'})
1016 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='baz')), objA)
1017 self.assertEqual(butler.get('foo', dp.DataId({'bar': 2}, tag='baz')), objB)
1018 del butler
1020 # create a New Butler repoC and put objC in it:
1021 objC = tstObj('c')
1022 repoCDir = os.path.join(self.testDir, 'repoC')
1023 butler = dp.Butler(inputs=self.repoBDir, outputs=repoCDir)
1024 butler.put(objC, 'foo', {'bar': 3})
1025 del butler
1027 # verify that repoC can be used as an input, and objA, objB, and objC can be gotten from it:
1028 butler = dp.Butler(inputs=repoCDir)
1029 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1})), objA)
1030 self.assertEqual(butler.get('foo', dp.DataId({'bar': 2})), objB)
1031 self.assertEqual(butler.get('foo', dp.DataId({'bar': 3})), objC)
1032 del butler
1034 # apply a tag and verify that repoC can be used as an input, and objA, objB, and objC can be gotten
1035 # from it:
1036 butler = dp.Butler(inputs={'root': repoCDir, 'tags': 'baz'})
1037 self.assertEqual(butler.get('foo', dp.DataId({'bar': 1}, tag='baz')), objA)
1038 self.assertEqual(butler.get('foo', dp.DataId({'bar': 2}, tag='baz')), objB)
1039 self.assertEqual(butler.get('foo', dp.DataId({'bar': 3}, tag='baz')), objC)
1040 del butler
1043class MemoryTester(lsst.utils.tests.MemoryTestCase):
1044 pass
1047if __name__ == '__main__': 1047 ↛ 1048line 1047 didn't jump to line 1048, because the condition on line 1047 was never true
1048 lsst.utils.tests.init()
1049 unittest.main()