From 8f66375e961003f291c3777976530a728d1d8b12 Mon Sep 17 00:00:00 2001
From: Andreas Linde <mail@andreaslinde.de>
Date: Mon, 5 Feb 2024 19:08:51 +0100
Subject: [PATCH] Fix heartbeats not being sent

---
 spine/entity_local.go  | 5 -----
 spine/feature_local.go | 8 ++++++++
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/spine/entity_local.go b/spine/entity_local.go
index 44d3ffb..a1eb351 100644
--- a/spine/entity_local.go
+++ b/spine/entity_local.go
@@ -66,11 +66,6 @@ func (r *EntityLocal) GetOrAddFeature(featureType model.FeatureTypeType, role mo
 	f.SetDescriptionString(description)
 	r.features = append(r.features, f)
 
-	if role == model.RoleTypeServer && featureType == model.FeatureTypeTypeDeviceDiagnosis {
-		// Update HeartbeatManager
-		r.device.HeartbeatManager().SetLocalFeature(r, f)
-	}
-
 	return f
 }
 
diff --git a/spine/feature_local.go b/spine/feature_local.go
index f291e18..12e689e 100644
--- a/spine/feature_local.go
+++ b/spine/feature_local.go
@@ -69,6 +69,14 @@ func (r *FeatureLocal) AddFunctionType(function model.FunctionType, read, write
 		return
 	}
 	r.operations[function] = NewOperations(read, write)
+
+	if r.role == model.RoleTypeServer &&
+		r.ftype == model.FeatureTypeTypeDeviceDiagnosis &&
+		function == model.FunctionTypeDeviceDiagnosisHeartbeatData {
+		// Update HeartbeatManager
+		r.Device().HeartbeatManager().SetLocalFeature(r.Entity(), r)
+	}
+
 }
 
 func (r *FeatureLocal) AddResultHandler(handler api.FeatureResultInterface) {