refactor(tracker,sensor): combine registration and update sensor structs

joshuar · Nov 28, 2023 · 3aa8a79 · 3aa8a79
1 parent a5ff5c6
commit 3aa8a79
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 156 deletions.
diff --git a/internal/hass/sensor/sensor.go b/internal/hass/sensor/sensor.go
@@ -19,31 +19,11 @@ const (
 // SensorRegistrationInfo is the JSON structure required to register a sensor
 // with HA.
 type SensorRegistrationInfo struct {
-	State             interface{} `json:"state"`
-	StateAttributes   interface{} `json:"attributes,omitempty"`
-	UniqueID          string      `json:"unique_id"`
-	Type              string      `json:"type"`
-	Name              string      `json:"name"`
-	UnitOfMeasurement string      `json:"unit_of_measurement,omitempty"`
-	StateClass        string      `json:"state_class,omitempty"`
-	EntityCategory    string      `json:"entity_category,omitempty"`
-	Icon              string      `json:"icon,omitempty"`
-	DeviceClass       string      `json:"device_class,omitempty"`
-	Disabled          bool        `json:"disabled,omitempty"`
-}
-
-func (reg *SensorRegistrationInfo) RequestType() api.RequestType {
-	return api.RequestTypeRegisterSensor
-}
-
-func (reg *SensorRegistrationInfo) RequestData() json.RawMessage {
-	data, err := json.Marshal(reg)
-	if err != nil {
-		log.Debug().Err(err).
-			Msg("Unable to marshal sensor to json.")
-		return nil
-	}
-	return data
+	Name              string `json:"name,omitempty"`
+	UnitOfMeasurement string `json:"unit_of_measurement,omitempty"`
+	StateClass        string `json:"state_class,omitempty"`
+	EntityCategory    string `json:"entity_category,omitempty"`
+	DeviceClass       string `json:"device_class,omitempty"`
 }
 
 // SensorUpdateInfo is the JSON structure required to update HA with the current
@@ -56,12 +36,23 @@ type SensorUpdateInfo struct {
 	UniqueID        string      `json:"unique_id"`
 }
 
-func (upd *SensorUpdateInfo) RequestType() api.RequestType {
-	return api.RequestTypeUpdateSensorStates
+type SensorState struct {
+	SensorUpdateInfo
+	SensorRegistrationInfo
+	Disabled   bool `json:"disabled,omitempty"`
+	Registered bool `json:"-"`
+}
+
+func (s *SensorState) RequestType() api.RequestType {
+	if s.Registered {
+		return api.RequestTypeUpdateSensorStates
+	} else {
+		return api.RequestTypeRegisterSensor
+	}
 }
 
-func (upd *SensorUpdateInfo) RequestData() json.RawMessage {
-	data, err := json.Marshal(upd)
+func (s *SensorState) RequestData() json.RawMessage {
+	data, err := json.Marshal(s)
 	if err != nil {
 		log.Debug().Err(err).
 			Msg("Unable to marshal sensor to json.")

diff --git a/internal/tracker/sensor.go b/internal/tracker/sensor.go
@@ -39,30 +39,23 @@ func prettyPrintState(s Sensor) string {
 	return b.String()
 }
 
-func marshalSensorUpdate(s Sensor) *sensor.SensorUpdateInfo {
-	return &sensor.SensorUpdateInfo{
-		StateAttributes: s.Attributes(),
-		Icon:            s.Icon(),
-		State:           s.State(),
-		Type:            marshalClass(s.SensorType()),
-		UniqueID:        s.ID(),
-	}
-}
-
-func marshalSensorRegistration(s Sensor) *sensor.SensorRegistrationInfo {
-	return &sensor.SensorRegistrationInfo{
-		StateAttributes:   s.Attributes(),
-		DeviceClass:       marshalClass(s.DeviceClass()),
-		Icon:              s.Icon(),
-		Name:              s.Name(),
-		State:             s.State(),
-		Type:              marshalClass(s.SensorType()),
-		UniqueID:          s.ID(),
-		UnitOfMeasurement: s.Units(),
-		StateClass:        marshalClass(s.StateClass()),
-		EntityCategory:    s.Category(),
-		Disabled:          false,
+func marshallSensorState(state Sensor, registered bool) *sensor.SensorState {
+	s := &sensor.SensorState{}
+	s.StateAttributes = state.Attributes()
+	s.Icon = state.Icon()
+	s.State = state.State()
+	s.Type = marshalClass(state.SensorType())
+	s.UniqueID = state.ID()
+	s.Registered = registered
+	if !s.Registered {
+		s.Name = state.Name()
+		s.DeviceClass = marshalClass(state.DeviceClass())
+		s.StateClass = marshalClass(state.StateClass())
+		s.UnitOfMeasurement = state.Units()
+		s.EntityCategory = state.Category()
+		s.Disabled = false
 	}
+	return s
 }
 
 type ComparableStringer interface {

diff --git a/internal/tracker/sensor_test.go b/internal/tracker/sensor_test.go
diff --git a/internal/tracker/tracker.go b/internal/tracker/tracker.go
@@ -88,12 +88,7 @@ func (t *SensorTracker) send(ctx context.Context, sensorUpdate Sensor) {
 		return
 	}
 	registered := <-t.registry.IsRegistered(sensorUpdate.ID())
-	switch registered {
-	case true:
-		req = marshalSensorUpdate(sensorUpdate)
-	case false:
-		req = marshalSensorRegistration(sensorUpdate)
-	}
+	req = marshallSensorState(sensorUpdate, registered)
 	responseCh := make(chan interface{}, 1)
 	go api.ExecuteRequest(ctx, req, responseCh)
 	response := <-responseCh