Update HeartbeatManager

Instead of checking for the proper feature entity and feature, use a setter

spine/api.go

@@ -167,6 +167,7 @@ type SubscriptionManager interface {
 type HeartbeatManager interface {
 	IsHeartbeatRunning() bool
 	UpdateHeartbeatOnSubscriptions()
+	SetLocalFeature(entity *EntityLocalImpl, feature FeatureLocal)
 	StartHeartbeat() error
 	StopHeartbeat()
 }

spine/entity_local.go

@@ -49,6 +49,12 @@ func (r *EntityLocalImpl) GetOrAddFeature(featureType model.FeatureTypeType, rol
 	}
 	f.SetDescriptionString(description)
 	r.features = append(r.features, f)
+
+	if role == model.RoleTypeServer && featureType == model.FeatureTypeTypeDeviceDiagnosis {
+		// Update HeartbeatManagerImpl
+		r.device.HeartbeatManager().SetLocalFeature(r, f)
+	}
+
 	return f
 }
 

spine/heartbeat_manager.go

@@ -47,7 +47,7 @@ func (c *HeartbeatManagerImpl) IsHeartbeatRunning() bool {
 // check if there are any heartbeat subscriptions left, otherwise stop creating new ones
 // or start creating heartbeats again if needed
 func (c *HeartbeatManagerImpl) UpdateHeartbeatOnSubscriptions() {
-	if err := c.updateLocal(); err != nil {
+	if c.localEntity == nil {
 		return
 	}
 
@@ -66,33 +66,18 @@ func (c *HeartbeatManagerImpl) UpdateHeartbeatOnSubscriptions() {
 	}
 }
 
-func (c *HeartbeatManagerImpl) updateLocal() error {
-	if c.localFeature == nil {
-		// search through all entities on the device
-		entities := c.localDevice.Entities()
-		for _, entity := range entities {
-			localFeature := entity.FeatureOfTypeAndRole(model.FeatureTypeTypeDeviceDiagnosis, model.RoleTypeServer)
-			if localFeature != nil {
-				c.localFeature = localFeature
-				c.localEntity = entity
-			}
-		}
-
-		if c.localFeature == nil {
-			return errors.New("Local feature for heartbeat sender address could not be found")
-		}
-	}
-
-	return nil
+func (c *HeartbeatManagerImpl) SetLocalFeature(entity *EntityLocalImpl, feature FeatureLocal) {
+	c.localEntity = entity
+	c.localFeature = feature
 }
 
 // Start setting heartbeat data
 // Make sure the a required FeatureTypeTypeDeviceDiagnosis with the role server is present
 // otherwise this will end with an error
 // Note: Remote features need to have a subscription to get notifications
 func (c *HeartbeatManagerImpl) StartHeartbeat() error {
-	if err := c.updateLocal(); err != nil {
-		return err
+	if c.localEntity == nil {
+		return errors.New("unknown entity")
 	}
 
 	timeout, err := c.heartBeatTimeout.GetTimeDuration()