From 29b3d39e163d5ddfd4c710cde34bb036497b991b Mon Sep 17 00:00:00 2001
From: Austin <7528924+Austinpayne@users.noreply.github.com>
Date: Fri, 17 May 2024 12:47:01 -0600
Subject: [PATCH] Support passing arguments to async main (#568)

* Support passing arguments to async main

* Add test for AsyncParsableCommand.main()
---
 .../Parsable Types/AsyncParsableCommand.swift | 30 ++++++--
 .../AsyncCommandEndToEndTests.swift           | 69 +++++++++++++++++++
 .../CMakeLists.txt                            |  1 +
 3 files changed, 94 insertions(+), 6 deletions(-)
 create mode 100644 Tests/ArgumentParserEndToEndTests/AsyncCommandEndToEndTests.swift

diff --git a/Sources/ArgumentParser/Parsable Types/AsyncParsableCommand.swift b/Sources/ArgumentParser/Parsable Types/AsyncParsableCommand.swift
index 3b0349765..590431ed2 100644
--- a/Sources/ArgumentParser/Parsable Types/AsyncParsableCommand.swift	
+++ b/Sources/ArgumentParser/Parsable Types/AsyncParsableCommand.swift	
@@ -25,14 +25,18 @@ public protocol AsyncParsableCommand: ParsableCommand {
 
 @available(macOS 10.15, macCatalyst 13, iOS 13, tvOS 13, watchOS 6, *)
 extension AsyncParsableCommand {
-  /// Executes this command, or one of its subcommands, with the program's
-  /// command-line arguments.
+  /// Executes this command, or one of its subcommands, with the given arguments.
   ///
-  /// Instead of calling this method directly, you can add `@main` to the root
-  /// command for your command-line tool.
-  public static func main() async {
+  /// This method parses an instance of this type, one of its subcommands, or
+  /// another built-in `AsyncParsableCommand` type, from command-line
+  /// (or provided) arguments, and then calls its `run()` method, exiting
+  /// with a relevant error message if necessary.
+  ///
+  /// - Parameter arguments: An array of arguments to use for parsing. If
+  ///   `arguments` is `nil`, this uses the program's command-line arguments.
+  public static func main(_ arguments: [String]?) async {
     do {
-      var command = try parseAsRoot()
+      var command = try parseAsRoot(arguments)
       if var asyncCommand = command as? AsyncParsableCommand {
         try await asyncCommand.run()
       } else {
@@ -42,6 +46,20 @@ extension AsyncParsableCommand {
       exit(withError: error)
     }
   }
+
+  /// Executes this command, or one of its subcommands, with the program's
+  /// command-line arguments.
+  ///
+  /// Instead of calling this method directly, you can add `@main` to the root
+  /// command for your command-line tool.
+  ///
+  /// This method parses an instance of this type, one of its subcommands, or
+  /// another built-in `AsyncParsableCommand` type, from command-line arguments,
+  /// and then calls its `run()` method, exiting with a relevant error message
+  /// if necessary.
+  public static func main() async {
+    await self.main(nil)
+  }
 }
 
 /// A type that can designate an `AsyncParsableCommand` as the program's
diff --git a/Tests/ArgumentParserEndToEndTests/AsyncCommandEndToEndTests.swift b/Tests/ArgumentParserEndToEndTests/AsyncCommandEndToEndTests.swift
new file mode 100644
index 000000000..52397ffa6
--- /dev/null
+++ b/Tests/ArgumentParserEndToEndTests/AsyncCommandEndToEndTests.swift
@@ -0,0 +1,69 @@
+//===----------------------------------------------------------*- swift -*-===//
+//
+// This source file is part of the Swift Argument Parser open source project
+//
+// Copyright (c) 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+//
+//===----------------------------------------------------------------------===//
+
+import XCTest
+import ArgumentParser
+
+final class AsyncCommandEndToEndTests: XCTestCase {
+}
+
+actor AsyncStatusCheck {
+  struct Status: OptionSet {
+    var rawValue: UInt8
+    
+    static var root: Self { .init(rawValue: 1 << 0) }
+    static var sub: Self  { .init(rawValue: 1 << 1) }
+  }
+  
+  @MainActor
+  var status: Status = []
+  
+  @MainActor
+  func update(_ status: Status) {
+    self.status.insert(status)
+  }
+}
+
+var statusCheck = AsyncStatusCheck()
+
+// MARK: AsyncParsableCommand.main() testing
+
+struct AsyncCommand: AsyncParsableCommand {
+  static var configuration: CommandConfiguration {
+    .init(subcommands: [SubCommand.self])
+  }
+  
+  func run() async throws {
+    await statusCheck.update(.root)
+  }
+  
+  struct SubCommand: AsyncParsableCommand {
+    func run() async throws {
+      await statusCheck.update(.sub)
+    }
+  }
+}
+
+extension AsyncCommandEndToEndTests {
+  @MainActor
+  func testAsyncMain_root() async throws {
+    XCTAssertFalse(statusCheck.status.contains(.root))
+    await AsyncCommand.main([])
+    XCTAssertTrue(statusCheck.status.contains(.root))
+  }
+  
+  @MainActor
+  func testAsyncMain_sub() async throws {
+    XCTAssertFalse(statusCheck.status.contains(.sub))
+    await AsyncCommand.main(["sub-command"])
+    XCTAssertTrue(statusCheck.status.contains(.sub))
+  }
+}
diff --git a/Tests/ArgumentParserEndToEndTests/CMakeLists.txt b/Tests/ArgumentParserEndToEndTests/CMakeLists.txt
index 2e094848c..a9069a635 100644
--- a/Tests/ArgumentParserEndToEndTests/CMakeLists.txt
+++ b/Tests/ArgumentParserEndToEndTests/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_library(EndToEndTests
+  AsyncCommandEndToEndTests.swift
   CustomParsingEndToEndTests.swift
   DefaultsEndToEndTests.swift
   EnumEndToEndTests.swift