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 

14from collections.abc import Callable 

15from pathlib import Path 

16from typing import Any, ClassVar 

17 

18import astropy.io.fits 

19import numpy as np 

20import pydantic 

21 

22from lsst.images.fits import FitsOutputArchive 

23from lsst.images.serialization import ArchiveTree, InputArchive 

24 

25 

26class _TinyTree(ArchiveTree): 

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

28 

29 SCHEMA_NAME: ClassVar[str] = "test_fits_output_archive" 

30 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

31 MIN_READ_VERSION: ClassVar[int] = 1 

32 PUBLIC_TYPE: ClassVar[type] = object 

33 

34 def deserialize( 

35 self, archive: InputArchive[Any], **kwargs: Any 

36 ) -> _TinyTree: # pragma: no cover - never invoked 

37 raise NotImplementedError() 

38 

39 

40class _PointerTarget(pydantic.BaseModel): 

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

42 

43 data: dict[str, Any] | None = None 

44 

45 

46def _write_archive(body: Callable[[FitsOutputArchive], None], tmp_path: Path) -> list[tuple[str, int | None]]: 

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

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

49 """ 

50 filename = tmp_path / "test.fits" 

51 with FitsOutputArchive.open(filename) as archive: 

52 body(archive) 

53 archive.add_tree(_TinyTree()) 

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

55 return [ 

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

57 for hdu in hdu_list[1:] 

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

59 ] 

60 

61 

62def test_repeated_direct_names_get_increasing_extver(tmp_path: Path) -> None: 

63 """Verify repeated direct names get increasing EXTVER disambiguation.""" 

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

65 sources = [] 

66 

67 def body(archive: FitsOutputArchive) -> None: 

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

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

70 

71 keys = _write_archive(body, tmp_path) 

72 assert sources == ["fits:DATA", "fits:DATA,2"] 

73 assert keys == [("DATA", None), ("DATA", 2)] 

74 

75 

76def test_direct_and_pointer_target_names_do_not_collide(tmp_path: Path) -> None: 

77 """Verify a direct name and a pointer target's nested name do not 

78 collide. 

79 """ 

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

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

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

83 # disambiguation. 

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

85 sources = [] 

86 

87 def serializer(archive: FitsOutputArchive): 

88 ref = archive.add_array(array, name="data") 

89 sources.append(ref.source) 

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

91 

92 def body(archive: FitsOutputArchive): 

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

94 archive.serialize_pointer("psf", serializer, key="psf-key") # type: ignore[arg-type] 

95 

96 keys = _write_archive(body, tmp_path) 

97 assert sources == ["fits:DATA", "fits:/DATA"] 

98 assert keys == [("DATA", None), ("/DATA", None)]