From 9c7fa49b50238db29c545f456f8367bcef345cc5 Mon Sep 17 00:00:00 2001
From: Erinome <rand256@yandex.ru>
Date: Wed, 5 May 2021 12:53:30 +0300
Subject: [PATCH] homeassistant autoconf doesn't like custom fields to death

---
 lib/MqttClient.js | 45 ++++++++++++++++++++++++++++-----------------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/lib/MqttClient.js b/lib/MqttClient.js
index 113149f..0fdc76a 100644
--- a/lib/MqttClient.js
+++ b/lib/MqttClient.js
@@ -21,7 +21,8 @@ const CUSTOM_COMMANDS = {
 	LOAD_MAP: "load_map",
 	STORE_MAP: "store_map",
 	GET_DESTINATIONS: "get_destinations",
-	PLAY_SOUND: "play_sound"
+	PLAY_SOUND: "play_sound",
+	SET_WATER_GRADE: "set_water_grade"
 };
 
 //TODO: since this is also displayed in the UI it should be moved somewhere else
@@ -106,7 +107,6 @@ const MqttClient = function (valetudo) {
 	this.topics = {
 		command: this.topicPrefix + "/" + this.identifier + "/command",
 		set_fan_speed: this.topicPrefix + "/" + this.identifier + "/set_fan_speed",
-		set_water_grade: this.topicPrefix + "/" + this.identifier + "/set_water_grade",
 		send_command: this.topicPrefix + "/" + this.identifier + "/custom_command",
 		state: this.topicPrefix + "/" + this.identifier + "/state",
 		map_data: this.topicPrefix + "/" + this.identifier + "/map_data",
@@ -169,9 +169,7 @@ const MqttClient = function (valetudo) {
 				command_topic: this.topics.command,
 				state_topic: this.topics.state,
 				set_fan_speed_topic: this.topics.set_fan_speed,
-				set_water_grade_topic: this.topics.set_water_grade,
 				fan_speed_list: Object.keys(FAN_SPEEDS),
-				water_grade_list: Object.keys(WATER_GRADES),
 				send_command_topic: this.topics.send_command,
 				json_attributes_topic: this.topics.attributes
 			}
@@ -215,7 +213,6 @@ MqttClient.prototype.connect = function () {
 			this.client.subscribe([
 				this.topics.command,
 				this.topics.set_fan_speed,
-				this.topics.set_water_grade,
 				this.topics.send_command
 			], {qos:this.qos}, err => {
 				if (!err) {
@@ -237,9 +234,6 @@ MqttClient.prototype.connect = function () {
 				case this.topics.set_fan_speed:
 					this.handleFanSpeedRequest(message);
 					break;
-				case this.topics.set_water_grade:
-					this.handleWaterGradeRequest(message);
-					break;
 				case this.topics.command:
 					this.handleCommand(message);
 					break;
@@ -442,14 +436,6 @@ MqttClient.prototype.handleFanSpeedRequest = function (speed) {
 	}
 };
 
-MqttClient.prototype.handleWaterGradeRequest = function (grade) {
-	if (Object.keys(WATER_GRADES).includes(grade)) {
-		this.vacuum.setWaterGrade(WATER_GRADES[grade], (err, data) => this.publishCommandStatus("set_water_grade",data,err));
-	} else if (parseInt(grade)) {
-		this.vacuum.setWaterGrade(parseInt(grade), (err, data) => this.publishCommandStatus("set_water_grade",data,err));
-	}
-};
-
 /**
  * @param command {string}
  */
@@ -781,8 +767,33 @@ MqttClient.prototype.handleCustomCommand = function (message) {
 					this.publishCommandStatus(msg.command,null,"Can't play specified file or location.");
 				}
 				break;
+			/**
+			 * set a water grade value on devices supporting it:
+			 * {
+			 *   "command": "set_water_grade",
+			 *   "grade": "medium"
+			 * }
+			 * digital codes can also be used as a value:
+			 * {
+			 *   "command": "set_water_grade",
+			 *   "grade": "201"
+			 * }
+			 */
+			case CUSTOM_COMMANDS.SET_WATER_GRADE:
+				if (this.vacuum.features.water_usage_ctrl) {
+					if (Object.keys(WATER_GRADES).includes(msg.grade)) {
+						this.vacuum.setWaterGrade(WATER_GRADES[msg.grade], (err, data) => this.publishCommandStatus("set_water_grade",data,err));
+					} else if (parseInt(msg.grade)) {
+						this.vacuum.setWaterGrade(parseInt(msg.grade), (err, data) => this.publishCommandStatus("set_water_grade",data,err));
+					} else {
+						this.publishCommandStatus(msg.command,null,"Unsupported water grade value specified.");
+					}
+				} else {
+					this.publishCommandStatus(msg.command,null,"Setting water grade is not supported on this device.");
+				}
+				break;
 			default:
-				this.publishCommandStatus(msg.command,null,"Received invalid custom command:" + JSON.stringify(msg));
+				this.publishCommandStatus(msg.command,null,"Received invalid custom command: " + JSON.stringify(msg));
 		}
 	}
 };