From 2d5d132713d0701d02d5ff33e9918812d13d2a83 Mon Sep 17 00:00:00 2001
From: Rylan Polster <rslpolster@gmail.com>
Date: Mon, 2 Jan 2023 14:33:33 -0500
Subject: [PATCH] Include `*flight` block source in cask API

---
 .../cask/artifact/abstract_flight_block.rb    |  5 ++
 Library/Homebrew/cask/cask.rb                 |  8 +--
 Library/Homebrew/test/cask/cask_spec.rb       | 50 +++++++++++++++++++
 .../fixtures/cask/Casks/conditional-flight.rb | 21 ++++++++
 4 files changed, 78 insertions(+), 6 deletions(-)
 create mode 100644 Library/Homebrew/test/support/fixtures/cask/Casks/conditional-flight.rb

diff --git a/Library/Homebrew/cask/artifact/abstract_flight_block.rb b/Library/Homebrew/cask/artifact/abstract_flight_block.rb
index eccdb10cab..d46011c160 100644
--- a/Library/Homebrew/cask/artifact/abstract_flight_block.rb
+++ b/Library/Homebrew/cask/artifact/abstract_flight_block.rb
@@ -1,6 +1,7 @@
 # typed: true
 # frozen_string_literal: true
 
+require "method_source"
 require "cask/artifact/abstract_artifact"
 
 module Cask
@@ -36,6 +37,10 @@ module Cask
         directives.keys.map(&:to_s).join(", ")
       end
 
+      def to_h
+        directives.transform_values(&:source)
+      end
+
       private
 
       def class_for_dsl_key(dsl_key)
diff --git a/Library/Homebrew/cask/cask.rb b/Library/Homebrew/cask/cask.rb
index b50acc9223..2e00a20a70 100644
--- a/Library/Homebrew/cask/cask.rb
+++ b/Library/Homebrew/cask/cask.rb
@@ -288,13 +288,9 @@ module Cask
 
     def artifacts_list
       artifacts.map do |artifact|
-        key, value = if artifact.is_a? Artifact::AbstractFlightBlock
-          artifact.summarize
-        else
-          [artifact.class.dsl_key, to_h_gsubs(artifact.to_args)]
-        end
+        next artifact.to_h if artifact.is_a? Artifact::AbstractFlightBlock
 
-        { key => value }
+        { artifact.class.dsl_key => to_h_gsubs(artifact.to_args) }
       end
     end
 
diff --git a/Library/Homebrew/test/cask/cask_spec.rb b/Library/Homebrew/test/cask/cask_spec.rb
index a6c563d012..d2b9a28c4e 100644
--- a/Library/Homebrew/test/cask/cask_spec.rb
+++ b/Library/Homebrew/test/cask/cask_spec.rb
@@ -1,6 +1,8 @@
 # typed: false
 # frozen_string_literal: true
 
+require "method_source"
+
 describe Cask::Cask, :cask do
   let(:cask) { described_class.new("versioned-cask") }
 
@@ -265,6 +267,46 @@ describe Cask::Cask, :cask do
         }
       JSON
     }
+    let(:expected_flight_variations) {
+      <<~JSON
+        {
+          "arm64_big_sur": {
+            "artifacts": [
+              {
+                "preflight": "    preflight do\\n      puts \\"preflight on Big Sur\\"\\n    end\\n"
+              },
+              {
+                "uninstall_postflight": "    uninstall_postflight do\\n      puts \\"uninstall_postflight on Big Sur\\"\\n    end\\n"
+              }
+            ]
+          },
+          "big_sur": {
+            "artifacts": [
+              {
+                "preflight": "    preflight do\\n      puts \\"preflight on Big Sur\\"\\n    end\\n"
+              },
+              {
+                "uninstall_postflight": "    uninstall_postflight do\\n      puts \\"uninstall_postflight on Big Sur\\"\\n    end\\n"
+              }
+            ]
+          },
+          "catalina": {
+            "artifacts": [
+              {
+                "preflight": "    preflight do\\n      puts \\"preflight on Catalina or older\\"\\n    end\\n"
+              }
+            ]
+          },
+          "mojave": {
+            "artifacts": [
+              {
+                "preflight": "    preflight do\\n      puts \\"preflight on Catalina or older\\"\\n    end\\n"
+              }
+            ]
+          }
+        }
+      JSON
+    }
 
     before do
       # Use a more limited symbols list to shorten the variations hash
@@ -300,5 +342,13 @@ describe Cask::Cask, :cask do
       expect(h).to be_a(Hash)
       expect(JSON.pretty_generate(h["variations"])).to eq expected_sha256_variations.strip
     end
+
+    it "returns the correct variations hash for a cask with conditional flight blocks" do
+      c = Cask::CaskLoader.load("conditional-flight")
+      h = c.to_hash_with_variations
+
+      expect(h).to be_a(Hash)
+      expect(JSON.pretty_generate(h["variations"])).to eq expected_flight_variations.strip
+    end
   end
 end
diff --git a/Library/Homebrew/test/support/fixtures/cask/Casks/conditional-flight.rb b/Library/Homebrew/test/support/fixtures/cask/Casks/conditional-flight.rb
new file mode 100644
index 0000000000..ea49b82c11
--- /dev/null
+++ b/Library/Homebrew/test/support/fixtures/cask/Casks/conditional-flight.rb
@@ -0,0 +1,21 @@
+cask "conditional-flight" do
+  version "1.2.3"
+  sha256 "67cdb8a02803ef37fdbf7e0be205863172e41a561ca446cd84f0d7ab35a99d94"
+
+  url "file://#{TEST_FIXTURE_DIR}/cask/caffeine/#{platform}/#{version}/#{arch}.zip"
+  homepage "https://brew.sh/"
+
+  on_big_sur do
+    preflight do
+      puts "preflight on Big Sur"
+    end
+    uninstall_postflight do
+      puts "uninstall_postflight on Big Sur"
+    end
+  end
+  on_catalina :or_older do
+    preflight do
+      puts "preflight on Catalina or older"
+    end
+  end
+end