Updated the shell-init command to work with habitat and omnibus installations

ashiqueps · ashiqueps · commit 5561d8cbd5c4 · 2025-02-04T18:50:45.000+05:30
Signed-off-by: Ashique Saidalavi &lt;Ashique.saidalavi@progress.com&gt;
diff --git a/chef-cli.gemspec b/chef-cli.gemspec
@@ -43,7 +43,7 @@ Gem::Specification.new do |gem|
   gem.add_dependency "mixlib-shellout", ">= 2.0", "< 4.0"
   gem.add_dependency "ffi-yajl", ">= 1.0", "< 3.0"
   gem.add_dependency "minitar", "~> 0.6"
-  gem.add_dependency "chef", "= 18.5.0"
+  gem.add_dependency "chef", "~> 18.0"
   gem.add_dependency "solve", "< 5.0", "> 2.0"
   gem.add_dependency "addressable", ">= 2.3.5", "< 2.9"
   gem.add_dependency "cookbook-omnifetch", "~> 0.5"
diff --git a/lib/chef-cli/command/env.rb b/lib/chef-cli/command/env.rb
@@ -85,7 +85,8 @@ def gem_environment
       end
 
       def paths
-        omnibus_env["PATH"].split(File::PATH_SEPARATOR)
+        env = habitat_install? ? habitat_env : omnibus_env
+        env["PATH"].split(File::PATH_SEPARATOR)
       rescue OmnibusInstallNotFound
         ENV["PATH"].split(File::PATH_SEPARATOR)
       end
diff --git a/lib/chef-cli/command/exec.rb b/lib/chef-cli/command/exec.rb
@@ -30,7 +30,8 @@ class Exec < ChefCLI::Command::Base
       def run(params)
         # Set ENV directly on the "parent" process (us) before running #exec to
         # ensure the custom PATH is honored when finding the command to exec
-        omnibus_env.each { |var, value| ENV[var] = value }
+        env = habitat_install? ? habitat_env : omnibus_env
+        env.each { |var, value| ENV[var] = value }
         exec(*params)
         raise "Exec failed without an exception, your ruby is buggy" # should never get here
       end
diff --git a/lib/chef-cli/command/license.rb b/lib/chef-cli/command/license.rb
@@ -75,6 +75,8 @@ def run(params)
         else
           ChefCLI::Licensing::Base.send(remaining_args[0])
         end
+      rescue ChefLicensing::LicenseKeyFetcher::LicenseKeyNotFetchedError
+        ui.msg("License key not fetched. Please try again.")
       end
 
       def debug?
diff --git a/lib/chef-cli/command/shell_init.rb b/lib/chef-cli/command/shell_init.rb
@@ -29,6 +29,7 @@ module Mixlib
 module ChefCLI
 
   class ShellCompletionTemplateContext
+    include ChefCLI::Helpers
 
     def commands
       ChefCLI.commands_map.command_specs.inject({}) do |cmd_info, (_key, cmd_spec)|
@@ -40,6 +41,10 @@ def commands
     def get_binding
       binding
     end
+
+    def habitat?
+      habitat_install?
+    end
   end
 
   module Command
@@ -103,7 +108,7 @@ def run(argv)
           return 1
         end
 
-        env = habitat_env.dup
+        env = (habitat_install? ? habitat_env : omnibus_env).dup
         path = env.delete("PATH")
         export(shell_name, "PATH", path)
         env.each do |var_name, value|
diff --git a/lib/chef-cli/completions/bash.sh.erb b/lib/chef-cli/completions/bash.sh.erb
@@ -2,4 +2,4 @@ _chef_comp() {
     local COMMANDS="<%= commands.keys.join(' ')-%>"
     COMPREPLY=($(compgen -W "$COMMANDS" -- ${COMP_WORDS[COMP_CWORD]} ))
 }
-complete -F _chef_comp chef
+complete -F _chef_comp <%= habitat? ? "chef-cli" : "chef" %>
diff --git a/lib/chef-cli/completions/chef.fish.erb b/lib/chef-cli/completions/chef.fish.erb
@@ -5,5 +5,5 @@
 set -l chef_commands <%= commands.keys.join(' ') %>;
 
 <% commands.each do |command, desc| -%>
-complete -c chef -f -n "not __fish_seen_subcommand_from $chef_commands" -a <%= command %> -d "<%= desc %>";
+complete -c <%= habitat? ? "chef-cli" : "chef" %> -f -n "not __fish_seen_subcommand_from $chef_commands" -a <%= command %> -d "<%= desc %>";
 <% end -%>
diff --git a/lib/chef-cli/completions/zsh.zsh.erb b/lib/chef-cli/completions/zsh.zsh.erb
@@ -17,5 +17,5 @@ function _chef() {
   fi
 }
 
-compdef _chef chef
+compdef _chef <%= habitat? ? "chef-cli" : "chef" %>
 
diff --git a/lib/chef-cli/dist.rb b/lib/chef-cli/dist.rb
@@ -10,6 +10,9 @@ class Dist
     CLI_PRODUCT = "Chef CLI".freeze
     CLI_GEM = "chef-cli".freeze
 
+    CHEF_DKE_PKG_NAME = "chef/chef-development-kit-enterprise".freeze
+    HAB_PKG_NAME = "chef/chef-cli".freeze
+
     # the name of the overall infra product
     INFRA_PRODUCT = "Chef Infra".freeze
 
diff --git a/lib/chef-cli/helpers.rb b/lib/chef-cli/helpers.rb
@@ -60,6 +60,21 @@ def omnibus_install?
       File.exist?(expected_omnibus_root) && File.exist?(File.join(expected_omnibus_root, "version-manifest.json"))
     end
 
+    # The habitat version of the chef-cli can be installed with standalone or chef-development-kit-enterprise
+    # This method checks if the habitat version of chef-cli is installed as standalone
+    def habitat_standalone?
+      @hab_standalone ||= (hab_pkg_installed?(ChefCLI::Dist::HAB_PKG_NAME) && !habitat_chef_dke?)
+    end
+
+    # This method checks if the habitat version of chef-cli is installed with chef-development-kit-enterprise
+    def habitat_chef_dke?
+      @hab_dke ||= hab_pkg_installed?(ChefCLI::Dist::CHEF_DKE_PKG_NAME)
+    end
+
+    def habitat_install?
+      habitat_chef_dke? || habitat_standalone?
+    end
+
     def omnibus_root
       @omnibus_root ||= omnibus_expand_path(expected_omnibus_root)
     end
@@ -110,29 +125,30 @@ def git_windows_bin_dir
       @git_windows_bin_dir ||= File.expand_path(File.join(omnibus_root, "embedded", "git", "usr", "bin"))
     end
 
-    # #
-    # # environment vars for habitat
-    # #
-    # def habitat_env
-    #   @habitat_env ||=
-    #   begin
-    #     # Define the necessary paths for the Habitat environment
-    #     # Custom GEM_HOME within Habitat
-    #     pkg_prefix = get_pkg_prefix
-    #     vendor_dir = File.join(pkg_prefix, "vendor")
-    #     path = [
-    #       File.join(pkg_prefix, "bin"),
-    #       ENV["PATH"].split(File::PATH_SEPARATOR), # Preserve existing PATH
-    #     ].flatten.uniq
-
-    #     {
-    #     "PATH" => path.join(File::PATH_SEPARATOR),
-    #     "GEM_ROOT" => Gem.default_dir, # Default directory for gems
-    #     "GEM_HOME" => vendor_dir,  # GEM_HOME pointing to the vendor directory
-    #     "GEM_PATH" => vendor_dir,  # GEM_PATH also pointing to the vendor directory
-    #     }
-    #   end
-    # end
+    #
+    # environment vars for habitat
+    #
+    def habitat_env
+      @habitat_env ||=
+      begin
+        # Define the necessary paths for the Habitat environment
+        # If it is a chef-dke installation, we will use the chef-dke bin path.
+        # Otherwise, we will use the chef-cli bin path.
+        bin_pkg_prefix = get_pkg_prefix(habitat_chef_dke? ? ChefCLI::Dist::CHEF_DKE_PKG_NAME : ChefCLI::Dist::HAB_PKG_NAME)
+        vendor_dir = File.join(get_pkg_prefix(ChefCLI::Dist::HAB_PKG_NAME), "vendor")
+        path = [
+          File.join(bin_pkg_prefix, "bin"),
+          ENV["PATH"].split(File::PATH_SEPARATOR), # Preserve existing PATH
+        ].flatten.uniq
+
+        {
+        "PATH" => path.join(File::PATH_SEPARATOR),
+        "GEM_ROOT" => Gem.default_dir, # Default directory for gems
+        "GEM_HOME" => vendor_dir,  # GEM_HOME pointing to the vendor directory
+        "GEM_PATH" => vendor_dir,  # GEM_PATH also pointing to the vendor directory
+        }
+      end
+    end
 
     #
     # environment vars for omnibus
@@ -153,17 +169,15 @@ def omnibus_env
         end
     end
 
-    # def get_pkg_prefix
-    #   pkg_origin = "chef"
-    #   pkg_name = "#{pkg_origin}/chef-cli" # Your origin and package name
-    #   path = `hab pkg path #{pkg_name}`.strip
+    def get_pkg_prefix(pkg_name)
+      path = `hab pkg path #{pkg_name}`.strip
 
-    #   if $?.success? && !path.empty?
-    #     path
-    #   else
-    #     raise "Failed to get pkg_prefix for #{pkg_name}: #{path}"
-    #   end
-    # end
+      if $?.success? && !path.empty?
+        path
+      else
+        raise "Failed to get pkg_prefix for #{pkg_name}: #{path}"
+      end
+    end
 
     def omnibus_expand_path(*paths)
       dir = File.expand_path(File.join(paths))
@@ -209,5 +223,14 @@ def reset!
     def macos?
       !!(RUBY_PLATFORM =~ /darwin/)
     end
+
+    # @return [Boolean] Checks if a habitat package is installed.
+    # If habitat itself is not installed, this method will return false.
+    #
+    # @api private
+    #
+    def hab_pkg_installed?(pkg_name)
+      `hab pkg list #{pkg_name} 2>/dev/null`.include?(pkg_name) rescue false
+    end
   end
 end
diff --git a/lib/chef-cli/licensing/config.rb b/lib/chef-cli/licensing/config.rb
@@ -22,6 +22,5 @@
   config.chef_product_name = "workstation"
   config.chef_entitlement_id = "x6f3bc76-a94f-4b6c-bc97-4b7ed2b045c0"
   config.chef_executable_name = "chef"
-  # config.license_server_url = "https://services.chef.io/licensing"
-  config.license_server_url   = "https://licensing-acceptance.chef.co/License"
+  config.license_server_url = "https://services.chef.io/licensing"
 end
diff --git a/spec/unit/command/env_spec.rb b/spec/unit/command/env_spec.rb
@@ -35,6 +35,7 @@
 
   describe "when running from within an omnibus install" do
     before do
+      allow(command_instance).to receive(:habitat_install?).and_return false
       allow(command_instance).to receive(:omnibus_install?).and_return true
       allow(command_instance).to receive(:omnibus_embedded_bin_dir).and_return(omnibus_embedded_bin_dir)
       allow(command_instance).to receive(:omnibus_bin_dir).and_return(omnibus_bin_dir)
@@ -57,6 +58,7 @@
   end
   describe "when running locally" do
     before do
+      allow(command_instance).to receive(:habitat_install?).and_return false
       allow(command_instance).to receive(:omnibus_install?).and_return false
       command_instance.ui = ui
     end
diff --git a/spec/unit/command/exec_spec.rb b/spec/unit/command/exec_spec.rb
@@ -57,6 +57,7 @@ def run_command
     let(:ruby_path) { File.join(fixtures_path, "eg_omnibus_dir/valid/embedded/bin/ruby") }
 
     before do
+      allow(command_instance).to receive(:habitat_install?).and_return(false)
       allow(Gem).to receive(:ruby).and_return(ruby_path)
 
       # Using a fake path separator to keep to prevent people from accidentally
diff --git a/spec/unit/command/shell_init_spec.rb b/spec/unit/command/shell_init_spec.rb
@@ -39,6 +39,7 @@
 
     before do
       allow(::Dir).to receive(:exist?).and_call_original
+      allow(command_instance).to receive(:habitat_install?).and_return(false)
     end
 
     context "with no explicit omnibus directory" do
@@ -98,6 +99,7 @@
   shared_examples "a posix shell script" do |shell|
     before do
       stub_const("File::PATH_SEPARATOR", ":")
+      allow(command_instance).to receive(:habitat_install?).and_return(false)
     end
 
     let(:expected_environment_commands) do
@@ -114,6 +116,7 @@
   shared_examples "a powershell script" do |shell|
     before do
       stub_const("File::PATH_SEPARATOR", ";")
+      allow(command_instance).to receive(:habitat_install?).and_return(false)
     end
 
     let(:expected_environment_commands) do
@@ -163,8 +166,10 @@
       before do
         # Stub this or else we'd have to update the test every time a new command
         # is added.
+        allow(command_instance).to receive(:habitat_install?).and_return(false)
         allow(command_instance.shell_completion_template_context).to receive(:commands)
           .and_return(command_descriptions)
+        allow(command_instance.shell_completion_template_context).to receive(:habitat?).and_return(false)
 
         allow(command_instance).to receive(:omnibus_embedded_bin_dir).and_return(omnibus_embedded_bin_dir)
         allow(command_instance).to receive(:omnibus_bin_dir).and_return(omnibus_bin_dir)
@@ -228,6 +233,7 @@
       before do
         # Stub this or else we'd have to update the test every time a new command
         # is added.
+        allow(command_instance).to receive(:habitat_install?).and_return(false)
         allow(command_instance.shell_completion_template_context).to receive(:commands)
           .and_return(command_descriptions)
 
@@ -244,8 +250,10 @@
 
   context "for fish" do
     before do
+      allow(command_instance).to receive(:habitat_install?).and_return(false)
       stub_const("File::PATH_SEPARATOR", ":")
     end
+
     let(:expected_path) { [omnibus_bin_dir, user_bin_dir, omnibus_embedded_bin_dir, ENV["PATH"], git_bin_dir].join(":").split(":").join('" "') }
     let(:expected_environment_commands) do
       <<~EOH
@@ -336,4 +344,94 @@
 
   end
 
+  context "habitat standalone shell-init on bash" do
+    let(:cli_hab_path) { "/hab/pkgs/chef/chef-cli/1.0.0/123" }
+
+    let(:argv) { ["bash"] }
+
+    before do
+      allow(command_instance).to receive(:habitat_chef_dke?).and_return(false)
+      allow(command_instance).to receive(:habitat_standalone?).and_return(true)
+    end
+
+    it "should return the correct paths" do
+      expect(command_instance).to receive(:get_pkg_prefix).with("chef/chef-cli").twice.and_return(cli_hab_path)
+
+      command_instance.run(argv)
+      expect(stdout_io.string).to include("export PATH=\"#{cli_hab_path}/bin")
+      expect(stdout_io.string).to include("export GEM_HOME=\"#{cli_hab_path}/vendor")
+      expect(stdout_io.string).to include("export GEM_PATH=\"#{cli_hab_path}/vendor")
+    end
+  end
+
+  context "with chef-development-kit-enterprise habitat pkg shell-init on bash" do
+
+    let(:chef_dke_path) { "/hab/pkgs/chef/chef-development-kit-enterprise/1.0.0/123" }
+    let(:cli_hab_path) { "/hab/pkgs/chef/chef-cli/1.0.0/123" }
+
+    let(:argv) { ["bash"] }
+
+    before do
+      allow(command_instance).to receive(:habitat_chef_dke?).and_return(true)
+      allow(command_instance).to receive(:habitat_standalone?).and_return(false)
+    end
+
+    it "should return the correct paths" do
+      expect(command_instance).to receive(:get_pkg_prefix).with("chef/chef-development-kit-enterprise").and_return(chef_dke_path)
+      expect(command_instance).to receive(:get_pkg_prefix).with("chef/chef-cli").and_return(cli_hab_path)
+
+      command_instance.run(argv)
+      expect(stdout_io.string).to include("export PATH=\"#{chef_dke_path}/bin")
+      expect(stdout_io.string).to include("export GEM_HOME=\"#{cli_hab_path}/vendor")
+      expect(stdout_io.string).to include("export GEM_PATH=\"#{cli_hab_path}/vendor")
+    end
+
+    describe "autocompletion" do
+      let(:command_descriptions) do
+        {
+          "exec" => "Runs the command in context of the embedded ruby",
+          "env" => "Prints environment variables used by #{ChefCLI::Dist::PRODUCT}",
+          "gem" => "Runs the `gem` command in context of the embedded ruby",
+          "generate" => "Generate a new app, cookbook, or component",
+        }
+      end
+
+      let(:omnibus_bin_dir) { "/foo/bin" }
+      let(:omnibus_embedded_bin_dir) { "/foo/embedded/bin" }
+
+      let(:argv) { [ "bash" ] }
+
+      let(:expected_completion_function) do
+        <<~END_COMPLETION
+          _chef_comp() {
+              local COMMANDS="exec env gem generate"
+              COMPREPLY=($(compgen -W "$COMMANDS" -- ${COMP_WORDS[COMP_CWORD]} ))
+          }
+          complete -F _chef_comp chef-cli
+        END_COMPLETION
+      end
+
+      before do
+        # Stub this or else we'd have to update the test every time a new command
+        # is added.
+        allow(command_instance).to receive(:get_pkg_prefix).with("chef/chef-development-kit-enterprise").and_return(chef_dke_path)
+        allow(command_instance).to receive(:get_pkg_prefix).with("chef/chef-cli").and_return(cli_hab_path)
+        allow(command_instance.shell_completion_template_context).to receive(:commands)
+          .and_return(command_descriptions)
+        allow(command_instance.shell_completion_template_context).to receive(:habitat?).and_return(true)
+
+        allow(command_instance).to receive(:omnibus_embedded_bin_dir).and_return(omnibus_embedded_bin_dir)
+        allow(command_instance).to receive(:omnibus_bin_dir).and_return(omnibus_bin_dir)
+      end
+
+      it "generates a completion function for the chef command" do
+        command_instance.run(argv)
+        expect(stdout_io.string).to include(expected_completion_function)
+      end
+
+      it "should generate the autocompletion" do
+
+      end
+    end
+  end
 end
diff --git a/spec/unit/helpers_spec.rb b/spec/unit/helpers_spec.rb

Original file line number	Diff line number	Diff line change
`@@ -2,4 +2,4 @@ _chef_comp() {`
`2`	`2`	`local COMMANDS="<%= commands.keys.join(' ')-%>"`
`3`	`3`	`COMPREPLY=($(compgen -W "$COMMANDS" -- ${COMP_WORDS[COMP_CWORD]} ))`
`4`	`4`	`}`
`5`		`-complete -F _chef_comp chef`
	`5`	`+complete -F _chef_comp <%= habitat? ? "chef-cli" : "chef" %>`
Original file line number	Diff line number	Diff line change
`@@ -17,5 +17,5 @@ function _chef() {`
`17`	`17`	`fi`
`18`	`18`	`}`
`19`	`19`
`20`		`-compdef _chef chef`
	`20`	`+compdef _chef <%= habitat? ? "chef-cli" : "chef" %>`
`21`	`21`