Coverage for tests/test_fits_output_archive.py: 100%

1# This file is part of lsst-images. 

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# Use of this source code is governed by a 3-clause BSD-style 

10# license that can be found in the LICENSE file. 

11 

12from __future__ import annotations 

13 

14import os 

15import tempfile 

16import unittest 

17from typing import ClassVar 

18 

19import astropy.io.fits 

20import numpy as np 

21import pydantic 

22 

23from lsst.images.fits import FitsOutputArchive 

24from lsst.images.serialization import ArchiveTree 

25 

26 

27class _TinyTree(ArchiveTree): 

28 """Minimal concrete ArchiveTree for low-level archive writes.""" 

29 

30 SCHEMA_NAME: ClassVar[str] = "test_fits_output_archive" 

31 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

32 MIN_READ_VERSION: ClassVar[int] = 1 

33 PUBLIC_TYPE: ClassVar[type] = object 

34 

35 def deserialize(self, archive, **kwargs): # pragma: no cover - never invoked 

36 raise NotImplementedError() 

37 

38 

39class _PointerTarget(pydantic.BaseModel): 

40 """A trivial pointer-target model holding an array reference.""" 

41 

42 data: dict | None = None 

43 

44 

45class FitsOutputArchiveNameRegistryTestCase(unittest.TestCase): 

46 """Tests for repeated-name disambiguation in `FitsOutputArchive`.""" 

47 

48 def _write_archive(self, body) -> list[tuple[str, int | None]]: 

49 """Write an archive, applying ``body`` to it, and return the 

50 ``(EXTNAME, EXTVER)`` pairs of the resulting extension HDUs. 

51 """ 

52 with tempfile.TemporaryDirectory() as tmpdir: 

53 filename = os.path.join(tmpdir, "test.fits") 

54 with FitsOutputArchive.open(filename) as archive: 

55 body(archive) 

56 archive.add_tree(_TinyTree()) 

57 with astropy.io.fits.open(filename) as hdu_list: 

58 return [ 

59 (hdu.header["EXTNAME"], hdu.header.get("EXTVER")) 

60 for hdu in hdu_list[1:] 

61 if hdu.header.get("EXTNAME") not in ("JSON", "INDEX") 

62 ] 

63 

64 def test_repeated_direct_names_get_increasing_extver(self): 

65 array = np.zeros((2, 2), dtype=np.float32) 

66 sources = [] 

67 

68 def body(archive): 

69 sources.append(archive.add_array(array, name="data").source) 

70 sources.append(archive.add_array(array, name="data").source) 

71 

72 keys = self._write_archive(body) 

73 self.assertEqual(sources, ["fits:DATA", "fits:DATA,2"]) 

74 self.assertEqual(keys, [("DATA", None), ("DATA", 2)]) 

75 

76 def test_direct_and_pointer_target_names_do_not_collide(self): 

77 # A direct name and a pointer target's nested name (registered with 

78 # a leading slash because the pointer's nested archive is rooted at 

79 # "") already produce distinct EXTNAMEs, so neither needs EXTVER 

80 # disambiguation. 

81 array = np.zeros((2, 2), dtype=np.float32) 

82 sources = [] 

83 

84 def serializer(nested): 

85 ref = nested.add_array(array, name="data") 

86 sources.append(ref.source) 

87 return _PointerTarget(data=ref.model_dump()) 

88 

89 def body(archive): 

90 sources.append(archive.add_array(array, name="data").source) 

91 archive.serialize_pointer("psf", serializer, key="psf-key") 

92 

93 keys = self._write_archive(body) 

94 self.assertEqual(sources, ["fits:DATA", "fits:/DATA"]) 

95 self.assertEqual(keys, [("DATA", None), ("/DATA", None)])