From 83573dad14f134157a9afba92a69776bf594130f Mon Sep 17 00:00:00 2001
From: Jennifer Medina <jmedina@stsci.edu>
Date: Thu, 29 Jun 2023 15:25:19 -0400
Subject: [PATCH] added unit tests for validating cutouts and proc time

---
 astrocut/tests/test_cube_cut.py | 43 +++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/astrocut/tests/test_cube_cut.py b/astrocut/tests/test_cube_cut.py
index 7ba66dd6..339e668a 100644
--- a/astrocut/tests/test_cube_cut.py
+++ b/astrocut/tests/test_cube_cut.py
@@ -5,6 +5,7 @@
 import astropy.units as u
 import numpy as np
 import pytest
+import time
 from astropy import wcs
 from astropy.coordinates import SkyCoord
 from astropy.io import fits
@@ -527,3 +528,45 @@ def test_s3_cube_cut(tmp_path: Path):
     assert np.isclose(hdulist[1].data["FLUX_ERR"][200][1, 2], 1.1239403)
     assert hdulist[0].header["CAMERA"] == 2
     hdulist.close()
+
+
+@pytest.mark.parametrize("cutout_size", [[5, 10], 20, 31])
+def test_multiprocessing(cutout_size, tmp_path):
+
+    tmpdir = str(tmp_path)
+
+    coord = SkyCoord(217.42893801, -62.67949189, unit="deg", frame="icrs")
+    cube_file = "s3://stpubdata/tess/public/mast/tess-s0038-2-2-cube.fits"
+
+    cutf_0threads = CutoutFactory(threads=0)
+    start_0time = time.time()
+    cutout_0threads = cutf_0threads.cube_cut(cube_file, coordinates=coord,
+                                             output_path=tmpdir, verbose=False,
+                                             cutout_size=cutout_size)
+    time_0threads = time.time() - start_0time
+    
+    cutf_4threads = CutoutFactory(threads=4)
+    start_4time = time.time()
+    cutout_4threads = cutf_4threads.cube_cut(cube_file, coordinates=coord,
+                                             output_path=tmpdir, verbose=False,
+                                             cutout_size=cutout_size)
+    time_4threads = time.time() - start_4time
+    
+    cutf_8threads = CutoutFactory(threads=8)
+    start_8time = time.time()
+    cutout_8threads = cutf_8threads.cube_cut(cube_file, coordinates=coord,
+                                             output_path=tmpdir, verbose=False,
+                                             cutout_size=cutout_size)
+    time_8threads = time.time() - start_8time
+    
+    y, x = 1, 2
+    index = np.random.randint(0, len(fits.getdata(cutout_0threads)["FLUX"]) - 1)
+    pixels, means = [], []
+    for cutout in [cutout_0threads, cutout_4threads, cutout_8threads]:
+
+        pixels.append(fits.getdata(cutout)["FLUX"][index][y, x])
+        means.append(np.mean(fits.getdata(cutout)["FLUX"][index]))
+
+    assert len(set(pixels)) == 1, f"pixel values are different for coord (1, 2): {pixels}"
+    assert len(set(means)) == 1
+    assert time_0threads > time_4threads > time_8threads
\ No newline at end of file