From d42e79fd91a3474746d7f6ede9967e0856d04df2 Mon Sep 17 00:00:00 2001
From: Light_Quanta <zch22248@foxmail.com>
Date: Sat, 10 May 2025 23:27:04 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8DbeastCharge=E8=B0=83=E7=94=A8?=
 =?UTF-8?q?=E9=94=99=E8=AF=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../superbwarfare/entity/DPSGeneratorEntity.java     | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java b/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java
index e081a3531..aebc9f494 100644
--- a/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java
+++ b/src/main/java/com/atsuishio/superbwarfare/entity/DPSGeneratorEntity.java
@@ -348,12 +348,10 @@ public class DPSGeneratorEntity extends LivingEntity implements GeoEntity {
 
     public void beastCharge() {
         this.entityData.set(LEVEL, 7);
-        this.getEnergy().ifPresent(cap -> {
-            if (cap instanceof SyncedEntityEnergyStorage storage) {
-                storage.setCapacity(this.getMaxEnergy());
-                storage.setMaxExtract(this.getMaxTransfer());
-                storage.setEnergy(this.getMaxEnergy());
-            }
-        });
+        if (this.getCapability(Capabilities.EnergyStorage.ENTITY, null) instanceof SyncedEntityEnergyStorage storage) {
+            storage.setCapacity(this.getMaxEnergy());
+            storage.setMaxExtract(this.getMaxTransfer());
+            storage.setEnergy(this.getMaxEnergy());
+        }
     }
 }