From 15596a737422867fa938b2ae71a83aa5d6c0d03b Mon Sep 17 00:00:00 2001
From: Arnold Noronha <arnold@tdrhq.com>
Date: Fri, 20 Sep 2024 18:00:27 -0400
Subject: [PATCH] Add a mainBranch configuration for the plugin

Summary: See https://github.com/screenshotbot/screenshotbot-gradle-plugin/issues/2

Differential Revision: https://phabricator.tdrhq.com/D10062
---
 .idea/misc.xml                                        |  1 -
 .../gradle/plugin/AbstractIntegrationBuilder.java     |  1 +
 .../gradle/plugin/ScreenshotbotPlugin.java            | 11 +++++++++++
 .../gradle/plugin/UploadScreenshotsTask.java          |  6 ++++++
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/.idea/misc.xml b/.idea/misc.xml
index 49504ef..42c5278 100644
--- a/.idea/misc.xml
+++ b/.idea/misc.xml
@@ -1,4 +1,3 @@
-<?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="ExternalStorageConfigurationManager" enabled="true" />
   <component name="ProjectRootManager" version="2" languageLevel="JDK_17" default="true" project-jdk-name="17" project-jdk-type="JavaSDK" />
diff --git a/plugin/src/main/java/io/screenshotbot/gradle/plugin/AbstractIntegrationBuilder.java b/plugin/src/main/java/io/screenshotbot/gradle/plugin/AbstractIntegrationBuilder.java
index 9938002..5da7b8e 100644
--- a/plugin/src/main/java/io/screenshotbot/gradle/plugin/AbstractIntegrationBuilder.java
+++ b/plugin/src/main/java/io/screenshotbot/gradle/plugin/AbstractIntegrationBuilder.java
@@ -110,6 +110,7 @@ protected void prepareTask(Task task, Project project, String mode) {
                     it.mode = mode;
                     it.hostname = extension.getHostname();
                     it.batch = extension.getBatch();
+                    it.mainBranch = extension.getMainBranch();
 
                     it.mustRunAfter(inputTaskName);
                     it.doFirst((innerTask) -> {
diff --git a/plugin/src/main/java/io/screenshotbot/gradle/plugin/ScreenshotbotPlugin.java b/plugin/src/main/java/io/screenshotbot/gradle/plugin/ScreenshotbotPlugin.java
index 33b0ec2..4e9ae96 100644
--- a/plugin/src/main/java/io/screenshotbot/gradle/plugin/ScreenshotbotPlugin.java
+++ b/plugin/src/main/java/io/screenshotbot/gradle/plugin/ScreenshotbotPlugin.java
@@ -53,6 +53,17 @@ public String getBatch() {
         public void setBatch(String batch) {
             this.batch = batch;
         }
+
+        public String getMainBranch() {
+            return mainBranch;
+        }
+
+        public void setMainBranch(String mainBranch) {
+            this.mainBranch = mainBranch;
+        }
+
+        private String mainBranch;
+
     }
     @Override
     public void apply(Project target) {
diff --git a/plugin/src/main/java/io/screenshotbot/gradle/plugin/UploadScreenshotsTask.java b/plugin/src/main/java/io/screenshotbot/gradle/plugin/UploadScreenshotsTask.java
index 70ed263..efab31f 100644
--- a/plugin/src/main/java/io/screenshotbot/gradle/plugin/UploadScreenshotsTask.java
+++ b/plugin/src/main/java/io/screenshotbot/gradle/plugin/UploadScreenshotsTask.java
@@ -15,6 +15,8 @@ public class UploadScreenshotsTask extends BaseRecorderTask {
     public String mode = "record";
     public String batch = null;
 
+    public String mainBranch = null;
+
     @Inject
     public UploadScreenshotsTask(ExecOperations execOperations) {
         super(execOperations);
@@ -48,6 +50,10 @@ public void uploadChannel() {
                 args.add(this.batch);
             }
 
+            if (this.mainBranch != null && this.mainBranch.length() > 0) {
+                args.add("--main-branch");
+                args.add(this.mainBranch);
+            }
             it.setArgs(args);
         });
     }