1""" Test fits plugin functionality.
2"""
3import pytest
4from imageio.testing import run_tests_if_main, get_test_dir, need_internet
5
6import imageio
7from imageio.core import get_remote_file, Request, IS_PYPY
8
9test_dir = get_test_dir()
10
11try:
12    import astropy
13except ImportError:
14    astropy = None
15
16
17def setup_module():
18    # During this test, pretend that FITS is the default format
19    imageio.formats.sort("FITS")
20
21
22def teardown_module():
23    # Set back to normal
24    imageio.formats.sort()
25
26
27@pytest.mark.skipif("astropy is None")
28def test_fits_format():
29    need_internet()  # We keep the fits files in the imageio-binary repo
30
31    # Test selection
32    for name in ["fits", ".fits"]:
33        format = imageio.formats["fits"]
34        assert format.name == "FITS"
35        assert format.__module__.endswith(".fits")
36
37    # Test cannot read
38    png = get_remote_file("images/chelsea.png")
39    assert not format.can_read(Request(png, "ri"))
40    assert not format.can_write(Request(png, "wi"))
41
42
43@pytest.mark.skipif("astropy is None")
44def test_fits_reading():
45    """ Test reading fits """
46
47    need_internet()  # We keep the fits files in the imageio-binary repo
48
49    if IS_PYPY:
50        return  # no support for fits format :(
51
52    simple = get_remote_file("images/simple.fits")
53    multi = get_remote_file("images/multi.fits")
54    compressed = get_remote_file("images/compressed.fits.fz")
55
56    # One image
57    im = imageio.imread(simple)
58    ims = imageio.mimread(simple)
59    assert (im == ims[0]).all()
60    assert len(ims) == 1
61
62    # Multiple images
63    ims = imageio.mimread(multi)
64    assert len(ims) == 3
65
66    R = imageio.read(multi)
67    assert R.format.name == "FITS"
68    ims = list(R)  # == [im for im in R]
69    assert len(ims) == 3
70
71    # Fail
72    raises = pytest.raises
73    raises(IndexError, R.get_data, -1)
74    raises(IndexError, R.get_data, 3)
75    raises(RuntimeError, R.get_meta_data, None)  # no meta data support
76    raises(RuntimeError, R.get_meta_data, 0)  # no meta data support
77
78    # Compressed image
79    im = imageio.imread(compressed)
80    assert im.shape == (2042, 3054)
81
82
83run_tests_if_main()
84