Skip to content

Commit

Permalink
unit test, separate s3 to http logic into separate function
Browse files Browse the repository at this point in the history
  • Loading branch information
snbianco committed May 21, 2024
1 parent 62882f8 commit 8328638
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 15 deletions.
22 changes: 18 additions & 4 deletions astrocut/asdf_cutouts.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,22 @@
from astropy.modeling import models


def _get_cloud_http(s3_uri: str) -> str:
""" Get the HTTP URI of a cloud resource from an S3 URI
Parameters
----------
s3_uri : string
the S3 URI of the cloud resource
"""
# create file system
fs = s3fs.S3FileSystem(anon=True)

# open resource and get URL
with fs.open(s3_uri, 'rb') as f:
return f.url()


def get_center_pixel(gwcsobj: gwcs.wcs.WCS, ra: float, dec: float) -> tuple:
""" Get the center pixel from a roman 2d science image
Expand Down Expand Up @@ -248,12 +264,10 @@ def asdf_cut(input_file: str, ra: float, dec: float, cutout_size: int = 20,
an image cutout object
"""

# if file comes from AWS cloud bucket, get URL
# if file comes from AWS cloud bucket, get HTTP URL to open with asdf
file = input_file
if isinstance(input_file, str) and input_file.startswith('s3://'):
fs = s3fs.S3FileSystem()
with fs.open(input_file, 'rb') as f:
file = f.url()
file = _get_cloud_http(input_file)

# get the 2d image data
with asdf.open(file) as f:
Expand Down
42 changes: 31 additions & 11 deletions astrocut/tests/test_asdf_cut.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@

import pathlib
from unittest.mock import MagicMock, patch
import numpy as np
import pytest

Expand All @@ -12,7 +13,7 @@
from astropy.wcs.utils import pixel_to_skycoord
from gwcs import wcs
from gwcs import coordinate_frames as cf
from astrocut.asdf_cutouts import get_center_pixel, get_cutout, asdf_cut, _slice_gwcs
from astrocut.asdf_cutouts import get_center_pixel, get_cutout, asdf_cut, _slice_gwcs, _get_cloud_http


def make_wcs(xsize, ysize, ra=30., dec=45.):
Expand Down Expand Up @@ -99,16 +100,6 @@ def make_file(tmp_path, fakedata):
yield filename


def test_get_center_pixel(fakedata):
""" test we can get the correct center pixel """
# get the fake data
__, gwcs = fakedata

pixel_coordinates, wcs = get_center_pixel(gwcs, 30., 45.)
assert np.allclose(pixel_coordinates, (np.array(500.), np.array(500.)))
assert np.allclose(wcs.celestial.wcs.crval, np.array([30., 45.]))


@pytest.fixture()
def output(tmp_path):
""" fixture to create the output path """
Expand All @@ -121,6 +112,16 @@ def _output_file(ext='fits'):
yield _output_file


def test_get_center_pixel(fakedata):
""" test we can get the correct center pixel """
# get the fake data
__, gwcs = fakedata

pixel_coordinates, wcs = get_center_pixel(gwcs, 30., 45.)
assert np.allclose(pixel_coordinates, (np.array(500.), np.array(500.)))
assert np.allclose(wcs.celestial.wcs.crval, np.array([30., 45.]))


@pytest.mark.parametrize('quantity', [True, False], ids=['quantity', 'array'])
def test_get_cutout(output, fakedata, quantity):
""" test we can create a cutout """
Expand Down Expand Up @@ -312,3 +313,22 @@ def test_slice_gwcs(fakedata):
# gwcs footprint/bounding_box expects ((x0, x1), (y0, y1)) but cutout.bbox is in ((y0, y1), (x0, x1))
assert (gwcsobj.footprint(bounding_box=tuple(reversed(cutout.bbox_original))) == sliced.footprint()).all()


@patch('s3fs.S3FileSystem')
def test_get_cloud_http(mock_s3fs):
""" test we can get HTTP URI of cloud resource """
# mock s3 file system operations
HTTP_URI = "http_test"
mock_file = MagicMock()
mock_fs = MagicMock()
mock_file.url.return_value = HTTP_URI
mock_fs.open.return_value.__enter__.return_value = mock_file
mock_s3fs.return_value = mock_fs

s3_uri = "s3://test_bucket/test_file.asdf"
http_uri = _get_cloud_http(s3_uri)

assert http_uri == HTTP_URI
mock_s3fs.assert_called_once_with(anon=True)
mock_fs.open.assert_called_once_with(s3_uri, 'rb')
mock_file.url.assert_called_once()

0 comments on commit 8328638

Please sign in to comment.