diff --git a/homeassistant/components/integration/sensor.py b/homeassistant/components/integration/sensor.py
index b8e72c3be5c..cd3e376b792 100644
--- a/homeassistant/components/integration/sensor.py
+++ b/homeassistant/components/integration/sensor.py
@@ -110,16 +110,10 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
         self._method = integration_method
 
         self._name = name if name is not None else f"{source_entity} integral"
-
-        if unit_of_measurement is None:
-            self._unit_template = (
-                f"{'' if unit_prefix is None else unit_prefix}{{}}{unit_time}"
-            )
-            # we postpone the definition of unit_of_measurement to later
-            self._unit_of_measurement = None
-        else:
-            self._unit_of_measurement = unit_of_measurement
-
+        self._unit_template = (
+            f"{'' if unit_prefix is None else unit_prefix}{{}}{unit_time}"
+        )
+        self._unit_of_measurement = unit_of_measurement
         self._unit_prefix = UNIT_PREFIXES[unit_prefix]
         self._unit_time = UNIT_TIME[unit_time]
         self._attr_state_class = STATE_CLASS_TOTAL_INCREASING
@@ -135,10 +129,10 @@ class IntegrationSensor(RestoreEntity, SensorEntity):
                 _LOGGER.warning("Could not restore last state: %s", err)
             else:
                 self._attr_device_class = state.attributes.get(ATTR_DEVICE_CLASS)
-
-                self._unit_of_measurement = state.attributes.get(
-                    ATTR_UNIT_OF_MEASUREMENT
-                )
+                if self._unit_of_measurement is None:
+                    self._unit_of_measurement = state.attributes.get(
+                        ATTR_UNIT_OF_MEASUREMENT
+                    )
 
         @callback
         def calc_integration(event):