Coverage for tests/test_serialization_streams.py: 96%

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"""Tests for reading archives from in-memory bytes and binary streams.""" 

12 

13from __future__ import annotations 

14 

15import gzip 

16import io 

17import os 

18from collections.abc import Callable 

19from pathlib import Path 

20 

21import numpy as np 

22import pytest 

23 

24from lsst.images import Box, Image, Mask 

25from lsst.images.fits import FitsInputArchive 

26from lsst.images.json import JsonInputArchive 

27from lsst.images.serialization import open_archive, read_archive, write_archive 

28 

29try: 

30 import h5py # noqa: F401 -- detect availability for NDF stream skips 

31 

32 from lsst.images.ndf import NdfInputArchive 

33 

34 H5PY_AVAILABLE = True 

35except ImportError: 

36 H5PY_AVAILABLE = False 

37 

38try: 

39 from compression import zstd as _stdlib_zstd # noqa: F401 -- detect zstd availability 

40 

41 ZSTD_AVAILABLE = True 

42except ImportError: 

43 try: 

44 import zstandard # noqa: F401 

45 

46 ZSTD_AVAILABLE = True 

47 except ImportError: 

48 ZSTD_AVAILABLE = False 

49 

50skip_no_h5py = pytest.mark.skipif(not H5PY_AVAILABLE, reason="h5py is not installed") 

51skip_no_zstd = pytest.mark.skipif(not ZSTD_AVAILABLE, reason="no zstd decompressor available") 

52 

53LOCAL_DATA_DIR = os.path.join(os.path.dirname(__file__), "data", "schema_v1") 

54 

55 

56def _zstd_compress(data: bytes) -> bytes: 

57 """Compress with whichever zstd library is available.""" 

58 try: 

59 from compression import zstd 

60 except ImportError: 

61 import zstandard 

62 

63 return zstandard.ZstdCompressor().compress(data) 

64 return zstd.compress(data) 

65 

66 

67def _make_image() -> Image: 

68 """Return a small float32 Image for stream round-trip tests.""" 

69 return Image(np.arange(16, dtype=np.float32).reshape(4, 4), bbox=Box.factory[0:4, 0:4]) 

70 

71 

72def _serialized_bytes(obj: object, extension: str, tmp_path: Path) -> bytes: 

73 """Write ``obj`` to a file under ``tmp_path`` and return its content.""" 

74 path = tmp_path / f"x{extension}" 

75 write_archive(obj, path) 

76 return path.read_bytes() 

77 

78 

79def test_fits_open_tree_stream(tmp_path: Path) -> None: 

80 """Verify FitsInputArchive.open_tree accepts a seekable binary stream.""" 

81 image = _make_image() 

82 stream = io.BytesIO(_serialized_bytes(image, ".fits", tmp_path)) 

83 with FitsInputArchive.open_tree(stream) as (archive, tree, info): 

84 assert info.schema_name == "image" 

85 result = tree.deserialize(archive) 

86 assert isinstance(result, Image) 

87 np.testing.assert_array_equal(result.array, image.array) 

88 

89 

90def test_json_open_tree_stream(tmp_path: Path) -> None: 

91 """Verify JsonInputArchive.open_tree accepts a seekable binary stream.""" 

92 image = _make_image() 

93 stream = io.BytesIO(_serialized_bytes(image, ".json", tmp_path)) 

94 with JsonInputArchive.open_tree(stream) as (archive, tree, info): 

95 result = tree.deserialize(archive) 

96 assert isinstance(result, Image) 

97 np.testing.assert_array_equal(result.array, image.array) 

98 

99 

100@skip_no_h5py 

101def test_ndf_open_tree_stream(tmp_path: Path) -> None: 

102 """Verify NdfInputArchive.open_tree accepts a seekable binary stream.""" 

103 image = _make_image() 

104 stream = io.BytesIO(_serialized_bytes(image, ".sdf", tmp_path)) 

105 with NdfInputArchive.open_tree(stream) as (archive, tree, info): 

106 result = tree.deserialize(archive) 

107 assert isinstance(result, Image) 

108 np.testing.assert_array_equal(result.array, image.array) 

109 

110 

111def test_read_stream_fits(tmp_path: Path) -> None: 

112 """Verify read_archive() turns in-memory FITS bytes into an Image.""" 

113 image = _make_image() 

114 result = read_archive(io.BytesIO(_serialized_bytes(image, ".fits", tmp_path))) 

115 assert isinstance(result, Image) 

116 np.testing.assert_array_equal(result.array, image.array) 

117 

118 

119def test_read_stream_json(tmp_path: Path) -> None: 

120 """Verify read_archive() turns in-memory JSON bytes into an Image.""" 

121 image = _make_image() 

122 result = read_archive(io.BytesIO(_serialized_bytes(image, ".json", tmp_path))) 

123 assert isinstance(result, Image) 

124 np.testing.assert_array_equal(result.array, image.array) 

125 

126 

127@skip_no_h5py 

128def test_read_stream_ndf(tmp_path: Path) -> None: 

129 """Verify read_archive() turns in-memory NDF bytes into an Image.""" 

130 image = _make_image() 

131 result = read_archive(io.BytesIO(_serialized_bytes(image, ".sdf", tmp_path))) 

132 assert isinstance(result, Image) 

133 np.testing.assert_array_equal(result.array, image.array) 

134 

135 

136def test_read_stream_cls_match_and_mismatch(tmp_path: Path) -> None: 

137 """Verify stream reads honor cls= for both matching and mismatched 

138 types. 

139 """ 

140 data = _serialized_bytes(_make_image(), ".fits", tmp_path) 

141 result = read_archive(io.BytesIO(data), cls=Image) 

142 assert isinstance(result, Image) 

143 with pytest.raises(TypeError): 

144 read_archive(io.BytesIO(data), cls=Mask) 

145 

146 

147def test_read_stream_kwargs_forwarded(tmp_path: Path) -> None: 

148 """Verify stream reads forward deserialize kwargs like bbox.""" 

149 big = Image(np.arange(64, dtype=np.float32).reshape(8, 8), bbox=Box.factory[0:8, 0:8]) 

150 sub = read_archive(io.BytesIO(_serialized_bytes(big, ".fits", tmp_path)), bbox=Box.factory[2:6, 2:6]) 

151 assert sub.array.shape == (4, 4) 

152 np.testing.assert_array_equal(sub.array, big.array[2:6, 2:6]) 

153 

154 

155def test_read_stream_format_override(tmp_path: Path) -> None: 

156 """Verify format= forces the named backend for a stream read.""" 

157 image = _make_image() 

158 result = read_archive(io.BytesIO(_serialized_bytes(image, ".json", tmp_path)), format="json") 

159 assert isinstance(result, Image) 

160 

161 

162def test_read_stream_wrong_format_override(tmp_path: Path) -> None: 

163 """Verify FITS bytes forced through the JSON backend fail in that 

164 backend. 

165 """ 

166 with pytest.raises(ValueError): 

167 read_archive(io.BytesIO(_serialized_bytes(_make_image(), ".fits", tmp_path)), format="json") 

168 

169 

170def test_read_stream_unrecognized_bytes() -> None: 

171 """Verify unrecognized stream content raises ValueError naming known 

172 formats. 

173 """ 

174 with pytest.raises(ValueError) as exc_info: 

175 read_archive(io.BytesIO(b"certainly not a supported format")) 

176 assert "FITS" in str(exc_info.value) 

177 

178 

179def test_read_stream_bad_format_name(tmp_path: Path) -> None: 

180 """Verify an unknown format= name raises ValueError.""" 

181 with pytest.raises(ValueError): 

182 read_archive(io.BytesIO(_serialized_bytes(_make_image(), ".json", tmp_path)), format="asdf") 

183 

184 

185def test_read_compressed_stream_raises(tmp_path: Path) -> None: 

186 """Verify a compressed stream raises ValueError telling the caller to 

187 decompress. 

188 """ 

189 # Compressed streams are the caller's responsibility to decompress; 

190 # the sniff error says what to do. 

191 data = gzip.compress(_serialized_bytes(_make_image(), ".fits", tmp_path)) 

192 with pytest.raises(ValueError) as exc_info: 

193 read_archive(io.BytesIO(data)) 

194 assert "gzip" in str(exc_info.value) 

195 

196 

197def test_open_stream_components(tmp_path: Path) -> None: 

198 """Verify open_archive() accepts a stream and supports component 

199 reads. 

200 """ 

201 visit_image = read_archive(os.path.join(LOCAL_DATA_DIR, "visit_image.json")) 

202 stream = io.BytesIO(_serialized_bytes(visit_image, ".fits", tmp_path)) 

203 with open_archive(stream) as reader: 

204 assert reader.info.schema_name == "visit_image" 

205 assert isinstance(reader.metadata, dict) 

206 assert reader.get_component("sky_projection") is not None 

207 full = reader.read() 

208 assert type(full).__name__ == "VisitImage" 

209 

210 

211def _write_compressed( 

212 extension: str, compress: Callable[[bytes], bytes], tmp_path: Path 

213) -> tuple[Path, Image]: 

214 """Write a compressed serialized Image under ``tmp_path``. 

215 

216 Returns the compressed file's path and the image it contains. 

217 """ 

218 image = _make_image() 

219 data = compress(_serialized_bytes(image, extension, tmp_path)) 

220 suffix = ".gz" if compress is gzip.compress else ".zst" 

221 path = tmp_path / f"c{extension}{suffix}" 

222 path.write_bytes(data) 

223 return path, image 

224 

225 

226def test_read_compressed_path_fits_gz(tmp_path: Path) -> None: 

227 """Verify a .fits.gz path reads transparently through read_archive().""" 

228 # Regression: .fits.gz was dispatched to the FITS backend but 

229 # handed to it still compressed. 

230 path, image = _write_compressed(".fits", gzip.compress, tmp_path) 

231 result = read_archive(path) 

232 assert isinstance(result, Image) 

233 np.testing.assert_array_equal(result.array, image.array) 

234 

235 

236def test_read_compressed_path_json_gz(tmp_path: Path) -> None: 

237 """Verify a .json.gz path reads transparently through read_archive().""" 

238 path, _ = _write_compressed(".json", gzip.compress, tmp_path) 

239 assert isinstance(read_archive(path), Image) 

240 

241 

242@skip_no_zstd 

243def test_read_compressed_path_fits_zst(tmp_path: Path) -> None: 

244 """Verify a .fits.zst path reads transparently through read_archive().""" 

245 path, _ = _write_compressed(".fits", _zstd_compress, tmp_path) 

246 assert isinstance(read_archive(path), Image) 

247 

248 

249def test_open_compressed_path_fits_gz(tmp_path: Path) -> None: 

250 """Verify open_archive() reads a .fits.gz path transparently.""" 

251 path, _ = _write_compressed(".fits", gzip.compress, tmp_path) 

252 with open_archive(path) as reader: 

253 assert isinstance(reader.read(), Image)