handle jobs with no credit cost

ASFHyP3 · Feb 28, 2024 · f6a5c2d · f6a5c2d
1 parent b89b1ed
commit f6a5c2d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/src/hyp3_sdk/jobs.py b/src/hyp3_sdk/jobs.py
@@ -293,4 +293,4 @@ def filter_jobs(
         return Batch(filtered_jobs)
 
     def total_credit_cost(self):
-        return sum(job.credit_cost for job in self.jobs)
+        return sum(job.credit_cost for job in self.jobs if job.credit_cost is not None)
diff --git a/tests/test_jobs.py b/tests/test_jobs.py
@@ -446,6 +446,17 @@ def test_batch_total_credit_cost():
     batch = Batch()
     assert batch.total_credit_cost() == 0
 
+    batch = Batch([
+        Job.from_dict({
+            'job_type': 'foo',
+            'job_id': 'foo',
+            'request_time': '2024-01-01T00:00:00Z',
+            'status_code': 'foo',
+            'user_id': 'foo',
+        }),
+    ])
+    assert batch.total_credit_cost() == 0
+
     batch = Batch([
         Job.from_dict({
             'job_type': 'foo',
@@ -455,6 +466,13 @@ def test_batch_total_credit_cost():
             'user_id': 'foo',
             'credit_cost': 4
         }),
+        Job.from_dict({
+            'job_type': 'foo',
+            'job_id': 'foo',
+            'request_time': '2024-01-01T00:00:00Z',
+            'status_code': 'foo',
+            'user_id': 'foo',
+        }),
     ])
     assert batch.total_credit_cost() == 4