diff --git a/src/Controller/CustomerAppAPI/InvoiceController.php b/src/Controller/CustomerAppAPI/InvoiceController.php
index 2e9779eb..c4bc132f 100644
--- a/src/Controller/CustomerAppAPI/InvoiceController.php
+++ b/src/Controller/CustomerAppAPI/InvoiceController.php
@@ -14,6 +14,7 @@ use App\Ramcar\TransactionOrigin;
 use App\Entity\CustomerVehicle;
 use App\Entity\Promo;
 use App\Entity\Battery;
+use App\Entity\BatterySize;
 use App\Entity\Customer;
 use App\Entity\CustomerMetadata;
 
@@ -117,9 +118,9 @@ class InvoiceController extends ApiController
 
         // if we have a trade in, add it as well
         if (!empty($trade_in_type) && !empty($trade_in_batt)) {
-            $ti_batt_obj = $this->em->getRepository(Battery::class)->find($trade_in_batt);
-            if (!empty($ti_batt_obj)) {
-                $icrit->addEntry($ti_batt_obj, $trade_in_type, 1);
+            $ti_batt_size_obj = $this->em->getRepository(BatterySize::class)->find($trade_in_batt);
+            if (!empty($ti_batt_size_obj)) {
+                $icrit->addTradeInEntry($ti_batt_size_obj, $trade_in_type, 1);
             }
         }
 
diff --git a/src/Controller/CustomerAppAPI/JobOrderController.php b/src/Controller/CustomerAppAPI/JobOrderController.php
index bd8b8296..0082de25 100644
--- a/src/Controller/CustomerAppAPI/JobOrderController.php
+++ b/src/Controller/CustomerAppAPI/JobOrderController.php
@@ -35,6 +35,7 @@ use App\Ramcar\WarrantyClass;
 use App\Ramcar\HubCriteria;
 use App\Ramcar\DeliveryStatus;
 use App\Entity\Battery;
+use App\Entity\BatterySize;
 use App\Entity\Hub;
 use App\Entity\Promo;
 use App\Entity\JOEvent;
@@ -688,10 +689,9 @@ class JobOrderController extends ApiController
 
         // if we have a trade in, add it as well
         if (!empty($trade_in_type) && !empty($trade_in_batt)) {
-            $ti_batt_obj = $this->em->getRepository(Battery::class)->find($trade_in_batt);
-            if (!empty($ti_batt_obj)) {
-                $battery_size = $ti_batt_obj->getSize();
-                $icrit->addTradeInEntry($battery_size, $trade_in_type, 1);
+            $ti_batt_size_obj = $this->em->getRepository(BatterySize::class)->find($trade_in_batt);
+            if (!empty($ti_batt_size_obj)) {
+                $icrit->addTradeInEntry($ti_batt_size_obj, $trade_in_type, 1);
             }
         }