From e69c95f6d224c00d21883e1d8bc1406910c2c935 Mon Sep 17 00:00:00 2001
From: Spencer Nelson <spencer@b612foundation.org>
Date: Tue, 11 Jul 2023 13:32:01 -0700
Subject: [PATCH] Add propagation test

---
 adam_core/propagator/tests/test_propagator.py | 45 +++++++++++++++++++
 1 file changed, 45 insertions(+)
 create mode 100644 adam_core/propagator/tests/test_propagator.py

diff --git a/adam_core/propagator/tests/test_propagator.py b/adam_core/propagator/tests/test_propagator.py
new file mode 100644
index 00000000..7bd4bcc4
--- /dev/null
+++ b/adam_core/propagator/tests/test_propagator.py
@@ -0,0 +1,45 @@
+from astropy.time import Time
+from ...orbits import Orbits
+from ... import coordinates
+from ...utils.helpers.orbits import make_real_orbits
+from .. import Propagator
+
+import quivr
+import pyarrow as pa
+
+
+class MockPropagator(Propagator):
+    # MockPropagator propagates orbits by just setting the time of the orbits.
+    def _propagate_orbits(self, orbits: Orbits, times: Time) -> Orbits:
+        all_times = []
+        for t in times:
+            repeated_time = Time([t] * len(orbits))
+            orbits.coordinates.time = coordinates.Times.from_astropy(
+                repeated_time
+            ).to_structarray()
+            all_times.append(orbits)
+
+        return quivr.concatenate(all_times)
+
+    def _generate_ephemeris(self, orbits: Orbits, observers):
+        raise NotImplementedError("not impleemented")
+
+
+def test_propagator_single_worker():
+    orbits = make_real_orbits(10)
+    times = Time(["2020-01-01T00:00:00", "2020-01-01T00:00:01"])
+
+    prop = MockPropagator()
+    have = prop.propagate_orbits(orbits, times, num_jobs=1)
+
+    assert len(have) == len(orbits) * len(times)
+
+
+def test_propagator_multiple_workers():
+    orbits = make_real_orbits(10)
+    times = Time(["2020-01-01T00:00:00", "2020-01-01T00:00:01"])
+
+    prop = MockPropagator()
+    have = prop.propagate_orbits(orbits, times, num_jobs=4)
+
+    assert len(have) == len(orbits) * len(times)