Merge pull request #1302 from habitat-sh/fix-studio-enter

Approved by: @nobody from Nowhere
Merged by: The Sentinels

components/hab/src/command/studio.rs

@@ -23,6 +23,12 @@ use hcore::os::users;
 
 use config;
 use error::Result;
+use VERSION;
+
+pub const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
+const DOCKER_CMD: &'static str = "docker";
+const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
+const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
 
 pub fn start(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
     // If the `$HAB_ORIGIN` environment variable is not present, then see if a default is set in
@@ -59,6 +65,23 @@ pub fn start(ui: &mut UI, args: Vec<OsString>) -> Result<()> {
     inner::start(ui, args)
 }
 
+/// Retrieves the
+pub fn image_identifier() -> String {
+    let version: Vec<&str> = VERSION.split("/").collect();
+    henv::var(DOCKER_IMAGE_ENVVAR).unwrap_or(format!("{}:{}", DOCKER_IMAGE, version[0]))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use VERSION;
+
+    #[test]
+    fn retrieve_image_identifier() {
+        assert_eq!(image_identifier(), format!("{}:{}", DOCKER_IMAGE, VERSION));
+    }
+}
+
 #[cfg(target_os = "linux")]
 mod inner {
     use std::ffi::OsString;
@@ -111,20 +134,12 @@ mod inner {
     use hcore::env as henv;
     use hcore::fs::{CACHE_KEY_PATH, find_command};
 
-    use VERSION;
     use error::{Error, Result};
     use exec;
-
-    const DOCKER_CMD: &'static str = "docker";
-    const DOCKER_CMD_ENVVAR: &'static str = "HAB_DOCKER_BINARY";
-
-    const DOCKER_IMAGE: &'static str = "habitat-docker-registry.bintray.io/studio";
-    const DOCKER_IMAGE_ENVVAR: &'static str = "HAB_DOCKER_STUDIO_IMAGE";
+    use super::{DOCKER_CMD, DOCKER_CMD_ENVVAR, image_identifier};
 
     pub fn start(_ui: &mut UI, args: Vec<OsString>) -> Result<()> {
         let docker = henv::var(DOCKER_CMD_ENVVAR).unwrap_or(DOCKER_CMD.to_string());
-        let image = henv::var(DOCKER_IMAGE_ENVVAR)
-            .unwrap_or(format!("{}:{}", DOCKER_IMAGE, VERSION));
 
         let cmd = match find_command(&docker) {
             Some(cmd) => cmd,
@@ -155,7 +170,7 @@ mod inner {
             .into());
         cmd_args.push("--volume".into());
         cmd_args.push(format!("{}:/src", env::current_dir().unwrap().to_string_lossy()).into());
-        cmd_args.push(image.into());
+        cmd_args.push(image_identifier().into());
         cmd_args.extend_from_slice(args.as_slice());
 
         for var in vec!["http_proxy", "https_proxy"] {