From e740c08e073b3b3a3a90c8bf5109251e1823fa97 Mon Sep 17 00:00:00 2001
From: Michael Russell <michael.russell@elastic.co>
Date: Mon, 7 Jan 2019 16:28:00 +0100
Subject: [PATCH] Generate node roles instead of hard coding them

---
 elasticsearch/templates/statefulset.yaml  | 10 ++++------
 elasticsearch/tests/elasticsearch_test.py | 11 +++++++++++
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/elasticsearch/templates/statefulset.yaml b/elasticsearch/templates/statefulset.yaml
index 782fd447d..65c4e8291 100644
--- a/elasticsearch/templates/statefulset.yaml
+++ b/elasticsearch/templates/statefulset.yaml
@@ -154,12 +154,10 @@ spec:
             value: "{{ .Values.networkHost }}"
           - name: ES_JAVA_OPTS
             value: "{{ .Values.esJavaOpts }}"
-          - name: node.master
-            value: "{{ .Values.roles.master }}"
-          - name: node.data
-            value: "{{ .Values.roles.data }}"
-          - name: node.ingest
-            value: "{{ .Values.roles.ingest }}"
+          {{- range $role, $enabled := .Values.roles }}
+          - name: node.{{ $role }}
+            value: "{{ $enabled }}"
+          {{- end }}
 {{- if .Values.extraEnvs }}
 {{ toYaml .Values.extraEnvs | indent 10 }}
 {{- end }}
diff --git a/elasticsearch/tests/elasticsearch_test.py b/elasticsearch/tests/elasticsearch_test.py
index 4f46569d1..f1b39d833 100644
--- a/elasticsearch/tests/elasticsearch_test.py
+++ b/elasticsearch/tests/elasticsearch_test.py
@@ -233,6 +233,17 @@ def test_dont_set_initial_master_nodes_if_not_master_when_using_es_version_7():
     for e in env:
         assert e['name'] != 'cluster.initial_master_nodes'
 
+def test_enabling_machine_learning_role():
+    config = '''
+roles:
+  ml: "true"
+'''
+    r = helm_template(config)
+    env = r['statefulset'][uname]['spec']['template']['spec']['containers'][0]['env']
+
+    assert {'name': 'node.ml',
+            'value': 'true'} in env
+
 
 def test_adding_extra_env_vars():
     config = '''