easy-setyp: Adding DCOS_CLUSTER env variable (#994)

dcos/config.py

@@ -59,12 +59,24 @@ def get_attached_cluster_path():
     if not os.path.exists(path):
         return None
 
+    attached = None
     clusters = os.listdir(get_clusters_path())
+    cluster_envvar = os.environ.get(constants.DCOS_CLUSTER)
+
     for c in clusters:
         cluster_path = os.path.join(path, c)
+        name = get_config_val("cluster.name",
+                              load_from_path(
+                                  os.path.join(cluster_path, "dcos.toml")))
+        if cluster_envvar is not None \
+           and (cluster_envvar == c or cluster_envvar == name):
+            return cluster_path
         if os.path.exists(os.path.join(
                 cluster_path, constants.DCOS_CLUSTER_ATTACHED_FILE)):
-            return cluster_path
+            attached = cluster_path
+
+    if attached is not None:
+        return attached
 
     # if only one cluster, set as attached
     if len(clusters) == 1:

dcos/constants.py

@@ -8,6 +8,9 @@ DCOS_CLUSTERS_SUBDIR = "clusters"
 """Name of the subdirectory containing the configuration of all configured
 clusters"""
 
+DCOS_CLUSTER = 'DCOS_CLUSTER'
+"""Name of the environment variable pointing to the DC/OS cluster id."""
+
 DCOS_CLUSTER_ATTACHED_FILE = "attached"
 """Name of the file indicating the current attached cluster"""
 

tests/test_cluster.py

@@ -54,14 +54,19 @@ def test_set_attached():
         assert cluster.set_attached(cluster_path) is None
         assert config.get_attached_cluster_path() == cluster_path
 
-        assert cluster.set_attached(cluster_path) is None
-        assert config.get_attached_cluster_path() == cluster_path
-
         cluster_path2 = add_cluster_dir("b", tempdir)
         # attach cluster already attached
         assert cluster.set_attached(cluster_path2) is None
         assert config.get_attached_cluster_path() == cluster_path2
 
+        # attach cluster through environment
+        os.environ[constants.DCOS_CLUSTER] = "a"
+        assert config.get_attached_cluster_path() == cluster_path
+
+        # attach back to old cluster through environment
+        os.environ[constants.DCOS_CLUSTER] = "b"
+        assert config.get_attached_cluster_path() == cluster_path2
+
 
 @patch('dcos.http.get')
 def test_setup_cluster_config(mock_get):