From d01e46e6a76c610d950f0f5b6cf73e84a9415bf1 Mon Sep 17 00:00:00 2001
From: George Helman <gmhelman11@gmail.com>
Date: Sat, 8 Jun 2024 23:53:44 -0400
Subject: [PATCH] Dev review changes

---
 dear_petition/petition/api/viewsets.py        |  9 ++-
 .../petition/etl/tests/test_transform.py      | 70 +++++++++++++------
 2 files changed, 58 insertions(+), 21 deletions(-)

diff --git a/dear_petition/petition/api/viewsets.py b/dear_petition/petition/api/viewsets.py
index c0012a6b..0fb1b922 100644
--- a/dear_petition/petition/api/viewsets.py
+++ b/dear_petition/petition/api/viewsets.py
@@ -372,8 +372,15 @@ def combine_batches(self, request):
     )
     def assign_client_to_batch(self, request, pk):
         client_id = request.data['client_id']
+
+        try:
+            client = pm.Client.objects.get(pk=client_id)
+        except pm.Client.DoesNotExist:
+            return Response(
+                "Unknown client.", status=status.HTTP_400_BAD_REQUEST
+            )
         batch = self.get_object()
-        batch.client_id = client_id
+        batch.client = client
         batch.save()
         batch.adjust_for_new_client_dob()
         return Response({"batch_id": batch.pk})
diff --git a/dear_petition/petition/etl/tests/test_transform.py b/dear_petition/petition/etl/tests/test_transform.py
index 209d0496..bc40df75 100644
--- a/dear_petition/petition/etl/tests/test_transform.py
+++ b/dear_petition/petition/etl/tests/test_transform.py
@@ -49,35 +49,64 @@ def test_recalculate_petitions(petition):
     assert not petition.has_attachments()
 
 
+
+
 def test_combine_batches(batch, batch_file, fake_pdf):
     record = CIPRSRecordFactory(
         batch=batch, jurisdiction=constants.DISTRICT_COURT, county="DURHAM"
     )
-    offense = OffenseFactory(
-        disposition_method="PROBATION OTHER",
-        ciprs_record=record,
-        jurisdiction=constants.DISTRICT_COURT,
-    )
-    OffenseRecordFactory(offense=offense, action="CHARGED")
+    record.refresh_record_from_data()
 
     second_batch = BatchFactory()
+
+    second_record_data = {
+        "General": {"County": "DURHAM", "File No": "00GR000001"},
+        "Case Information": {
+            "Case Status": "DISPOSED",
+            "Offense Date": "2018-01-01T20:00:00",
+        },
+        "Defendant": {
+            "Date of Birth/Estimated Age": "1990-01-01",
+            "Name": "DOE,JON,BOJACK",
+            "Race": "WHITE",
+            "Sex": "MALE",
+        },
+        "District Court Offense Information": [
+            {
+                "Records": [
+                    {
+                        "Action": "CHARGED",
+                        "Description": "SPEEDING(80 mph in a 65 mph zone)",
+                        "Severity": "TRAFFIC",
+                        "Law": "20-141(J1)",
+                        "Code": "4450",
+                    },
+                ],
+                "Disposed On": "2018-02-01",
+                "Disposition Method": "DISPOSED BY JUDGE",
+            },
+            {
+                "Records": [
+                    {
+                        "Action": "CHARGED",
+                        "Description": "SPEEDING(80 mph in a 65 mph zone)",
+                        "Severity": "TRAFFIC",
+                        "Law": "20-141(J1)",
+                        "Code": "4450",
+                    },
+                ],
+                "Disposed On": "2018-02-01",
+                "Disposition Method": "DISPOSED BY JUDGE",
+            }
+        ],
+        "Superior Court Offense Information": [],
+    }
     second_batch_file = BatchFileFactory(batch=second_batch, file=fake_pdf)
 
     second_record = CIPRSRecordFactory(
-        batch=second_batch, batch_file=second_batch_file, jurisdiction=constants.DISTRICT_COURT, county="DURHAM"
-    )
-    second_offense = OffenseFactory(
-        disposition_method="PROBATION OTHER",
-        ciprs_record=second_record,
-        jurisdiction=constants.DISTRICT_COURT,
-    )
-    OffenseRecordFactory(offense=second_offense, action="CHARGED")
-    third_offense = OffenseFactory(
-        disposition_method="PROBATION OTHER",
-        ciprs_record=second_record,
-        jurisdiction=constants.SUPERIOR_COURT,
+        batch=second_batch, data = second_record_data, batch_file=second_batch_file, jurisdiction=constants.DISTRICT_COURT, county="DURHAM"
     )
-    OffenseRecordFactory(offense=third_offense, action="CHARGED")
+    second_record.refresh_record_from_data()
   
     assert batch.records.count() == 1
     assert pm.Offense.objects.filter(ciprs_record__batch__id=batch.id).count() == 1
@@ -89,7 +118,8 @@ def test_combine_batches(batch, batch_file, fake_pdf):
     new_batch = combine_batches([batch.id, second_batch.id], label=new_label, user_id=user_id)
 
     assert new_batch.records.count() == 2
-    assert pm.Offense.objects.filter(ciprs_record__batch__id=new_batch.id).count() == 2
+    assert pm.Offense.objects.filter(ciprs_record__batch__id=new_batch.id).count() == 3
+    assert pm.OffenseRecord.objects.filter(offense__ciprs_record__batch__id=new_batch.id).count() == 4
     assert new_batch.files.count() == 2
     assert new_batch.label == new_label
     assert new_batch.user_id == user_id