From f16238a36184dd1f871ba9b36bad891493911695 Mon Sep 17 00:00:00 2001
From: Ismael Mendoza <11745764+ismael-mendoza@users.noreply.github.com>
Date: Sat, 15 Oct 2022 16:43:27 -0400
Subject: [PATCH] add tests for rotation (#383)

---
 tests/test_draw.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/tests/test_draw.py b/tests/test_draw.py
index 676bcfc43..d085d3862 100644
--- a/tests/test_draw.py
+++ b/tests/test_draw.py
@@ -195,3 +195,28 @@ def test_shear_draw():
         apply_shear=True,
     )
     next(draw_generator)
+
+
+def test_rotation():
+    stamp_size = 24.0
+    seed = 0
+    catalog_name = "data/sample_input_catalog.fits"
+    catalog = btk.catalog.CatsimCatalog.from_file(catalog_name)
+    sampling_function_shear = btk.sampling_functions.DefaultSampling(
+        stamp_size=stamp_size, seed=seed, max_number=3, min_number=1
+    )
+    survey = btk.survey.get_surveys("LSST")
+    draw_generator = btk.draw_blends.CatsimGenerator(
+        catalog,
+        sampling_function_shear,
+        survey,
+        batch_size=10,
+        stamp_size=stamp_size,
+        cpus=1,
+        add_noise="all",
+        seed=seed,
+        augment_data=True,
+    )
+    blend_results = next(draw_generator)
+    rotation = blend_results["blend_list"][0]["btk_rotation"]
+    assert np.all(0 <= rotation) and np.all(rotation <= 360)