From 9256c6885db80fbbcb877d1b118da81e0e9799f7 Mon Sep 17 00:00:00 2001
From: Sabrina Yan <9669990+violetbrina@users.noreply.github.com>
Date: Fri, 13 Dec 2024 14:08:07 +1100
Subject: [PATCH] fix(cpg_flow_test/,cpg_flow_test/jobs/): fix dependency in
 first stage

---
 cpg_flow_test/jobs/first_n_primes.py      | 9 ++++-----
 cpg_flow_test/jobs/iterative_digit_sum.py | 2 --
 cpg_flow_test/stages.py                   | 2 +-
 3 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/cpg_flow_test/jobs/first_n_primes.py b/cpg_flow_test/jobs/first_n_primes.py
index 36c3272..aa93ea5 100644
--- a/cpg_flow_test/jobs/first_n_primes.py
+++ b/cpg_flow_test/jobs/first_n_primes.py
@@ -3,15 +3,14 @@
 from hailtop.batch.job import Job
 
 
-def first_n_primes(
-    b: Batch,
-    input_file_path: str,
-    output_file_path: str,
-) -> list[Job]:
+def first_n_primes(b: Batch, input_file_path: str, output_file_path: str, depends_on: Job) -> list[Job]:
     title = 'First N Primes'
     job = b.new_job(name=title)
     id_sum_path = b.read_input(input_file_path)
 
+    if depends_on:
+        job.depends_on(depends_on)
+
     cmd = f"""
     is_prime() {{
         local num=$1
diff --git a/cpg_flow_test/jobs/iterative_digit_sum.py b/cpg_flow_test/jobs/iterative_digit_sum.py
index 9eb53bd..eb93495 100644
--- a/cpg_flow_test/jobs/iterative_digit_sum.py
+++ b/cpg_flow_test/jobs/iterative_digit_sum.py
@@ -42,8 +42,6 @@ def iterative_digit_sum(
         iterative_digit_sum "$concatenated_digits"
         }}
 
-        # Example usage
-        # Replace "$1" with your actual input
         echo "Input: {sequencing_group.id}\n"
         result=$(extract_digits_and_sum {sequencing_group.id})
         echo "Result: $result\n"
diff --git a/cpg_flow_test/stages.py b/cpg_flow_test/stages.py
index 2c2a843..3878330 100644
--- a/cpg_flow_test/stages.py
+++ b/cpg_flow_test/stages.py
@@ -71,7 +71,7 @@ def queue_jobs(self, sequencing_group: SequencingGroup, inputs: StageInput) -> S
 
         # Generate first N primes
         primes_output_path = str(self.expected_outputs(sequencing_group).get('primes', ''))
-        job_primes = first_n_primes(b, id_sum_output_path, primes_output_path)
+        job_primes = first_n_primes(b, id_sum_output_path, primes_output_path, depends_on=job_id_sum)
 
         jobs = [job_id_sum, job_primes]