Update GHES generate script to accept --use-github-storage option

releasenotes/v1.8.1.md

@@ -1 +1,2 @@
-Add `--use-github-storage` to gh [gei|bbs2gh] migrate-repo command to support uploading to a GitHub owned storage
+Add `--use-github-storage` to `gh [gei|bbs2gh] migrate-repo` command to support uploading to a GitHub owned storage
+Add `--use-github-storage` to `gh gei generate-script` command to support uploading to a GitHub owned storage

src/OctoshiftCLI.IntegrationTests/GhesToGithub.cs

@@ -28,6 +28,7 @@ public sealed class GhesToGithub : IDisposable
     private readonly Dictionary<string, string> _tokens;
     private readonly DateTime _startTime;
     private readonly ArchiveUploader _archiveUploader;
+    private readonly string _azureStorageConnectionString;
 
     public GhesToGithub(ITestOutputHelper output)
     {
@@ -38,12 +39,11 @@ public GhesToGithub(ITestOutputHelper output)
 
         var sourceGithubToken = Environment.GetEnvironmentVariable("GHES_PAT");
         var targetGithubToken = Environment.GetEnvironmentVariable("GHEC_PAT");
-        var azureStorageConnectionString = Environment.GetEnvironmentVariable($"AZURE_STORAGE_CONNECTION_STRING_GHES_{TestHelper.GetOsName().ToUpper()}");
+        _azureStorageConnectionString = Environment.GetEnvironmentVariable($"AZURE_STORAGE_CONNECTION_STRING_GHES_{TestHelper.GetOsName().ToUpper()}");
         _tokens = new Dictionary<string, string>
         {
             ["GH_SOURCE_PAT"] = sourceGithubToken,
             ["GH_PAT"] = targetGithubToken,
-            ["AZURE_STORAGE_CONNECTION_STRING"] = azureStorageConnectionString
         };
 
         _versionClient = new HttpClient();
@@ -57,13 +57,16 @@ public GhesToGithub(ITestOutputHelper output)
         _targetGithubClient = new GithubClient(logger, _targetGithubHttpClient, new VersionChecker(_versionClient, logger), new RetryPolicy(logger), new DateTimeProvider(), targetGithubToken);
         _targetGithubApi = new GithubApi(_targetGithubClient, "https://api.github.com", new RetryPolicy(logger), _archiveUploader);
 
-        _blobServiceClient = new BlobServiceClient(azureStorageConnectionString);
+        _blobServiceClient = new BlobServiceClient(_azureStorageConnectionString);
 
         _sourceHelper = new TestHelper(_output, _sourceGithubApi, _sourceGithubClient) { GithubApiBaseUrl = GHES_API_URL };
         _targetHelper = new TestHelper(_output, _targetGithubApi, _targetGithubClient, _blobServiceClient);
     }
 
-    public async Task Basic()
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task Basic(bool useGithubStorage)
     {
         var githubSourceOrg = $"e2e-testing-{TestHelper.GetOsName()}";
         var githubTargetOrg = $"octoshift-e2e-ghes-{TestHelper.GetOsName()}";
@@ -72,9 +75,17 @@ public async Task Basic()
 
         var retryPolicy = new RetryPolicy(null);
 
+        if (!useGithubStorage)
+        {
+            _tokens["AZURE_STORAGE_CONNECTION_STRING"] = _azureStorageConnectionString;
+        }
+
         await retryPolicy.Retry(async () =>
         {
-            await _targetHelper.ResetBlobContainers();
+            if (!useGithubStorage)
+            {
+                await _targetHelper.ResetBlobContainers();
+            }
 
             await _sourceHelper.ResetGithubTestEnvironment(githubSourceOrg);
             await _targetHelper.ResetGithubTestEnvironment(githubTargetOrg);
@@ -83,8 +94,14 @@ await retryPolicy.Retry(async () =>
             await _sourceHelper.CreateGithubRepo(githubSourceOrg, repo2);
         });
 
-        await _targetHelper.RunGeiCliMigration(
-            $"generate-script --github-source-org {githubSourceOrg} --github-target-org {githubTargetOrg} --ghes-api-url {GHES_API_URL} --download-migration-logs", _tokens);
+        // Build the command with conditional option
+        var command = $"generate-script --github-source-org {githubSourceOrg} --github-target-org {githubTargetOrg} --ghes-api-url {GHES_API_URL} --download-migration-logs";
+        if (useGithubStorage)
+        {
+            command += " --use-github-storage";
+        }
+
+        await _targetHelper.RunGeiCliMigration(command, _tokens);
 
         _targetHelper.AssertNoErrorInLogs(_startTime);
 
@@ -96,6 +113,7 @@ await _targetHelper.RunGeiCliMigration(
         _targetHelper.AssertMigrationLogFileExists(githubTargetOrg, repo1);
         _targetHelper.AssertMigrationLogFileExists(githubTargetOrg, repo2);
     }
+
     public void Dispose()
     {
         _sourceGithubHttpClient?.Dispose();

src/OctoshiftCLI.Tests/gei/Commands/GenerateScript/GenerateScriptCommandArgsTests.cs

@@ -72,5 +72,40 @@ public void It_Throws_When_Invalid_URL_Provided()
                 .Should()
                 .Throw<OctoshiftCliException>();
         }
+
+        [Fact]
+        public void UseGithubStorage_Without_Ghes_Api_Url_Throws()
+        {
+            var args = new GenerateScriptCommandArgs
+            {
+                GithubSourceOrg = "foo",
+                GithubTargetOrg = "bar",
+                UseGithubStorage = true
+            };
+
+            FluentActions.Invoking(() => args.Validate(_mockOctoLogger.Object))
+                 .Should()
+                 .ThrowExactly<OctoshiftCliException>()
+                 .WithMessage("*--use-github-storage*");
+        }
+
+
+        [Fact]
+        public void UseGithubStorage_And_Aws_Bucket_Name_Throws()
+        {
+            var args = new GenerateScriptCommandArgs
+            {
+                GithubSourceOrg = "foo",
+                GithubTargetOrg = "bar",
+                AwsBucketName = "aws",
+                GhesApiUrl = "https://github.contoso.com",
+                UseGithubStorage = true
+            };
+
+            FluentActions.Invoking(() => args.Validate(_mockOctoLogger.Object))
+                 .Should()
+                 .ThrowExactly<OctoshiftCliException>()
+                 .WithMessage("*--use-github-storage flag was provided with an AWS S3 Bucket name*");
+        }
     }
 }

src/OctoshiftCLI.Tests/gei/Commands/GenerateScript/GenerateScriptCommandHandlerTests.cs

@@ -1102,6 +1102,68 @@ exit 1
             _script.Should().NotContain(expected);
         }
 
+        [Fact]
+        public async Task Sequential_Github_Single_Repo_With_UseGithubStorage()
+        {
+            // Arrange
+            var GHES_API_URL = "https://foo.com/api/v3";
+
+            _mockGithubApi
+                .Setup(m => m.GetRepos(SOURCE_ORG))
+                .ReturnsAsync(new[] { (REPO, "private") });
+
+            var expected = $"Exec {{ gh gei migrate-repo --github-source-org \"{SOURCE_ORG}\" --source-repo \"{REPO}\" --github-target-org \"{TARGET_ORG}\" --target-repo \"{REPO}\" --ghes-api-url \"{GHES_API_URL}\" --use-github-storage --target-repo-visibility private }}";
+
+            // Act
+            var args = new GenerateScriptCommandArgs
+            {
+                GithubSourceOrg = SOURCE_ORG,
+                GithubTargetOrg = TARGET_ORG,
+                Output = new FileInfo("unit-test-output"),
+                Sequential = true,
+                UseGithubStorage = true,
+                GhesApiUrl = GHES_API_URL,
+            };
+            await _handler.Handle(args);
+
+            _script = TrimNonExecutableLines(_script);
+
+            // Assert
+            _script.Should().Be(expected);
+        }
+
+        [Fact]
+        public async Task Parallel_Github_Single_Repo_With_UseGithubStorage()
+        {
+            // Arrange
+            _mockGithubApi
+                .Setup(m => m.GetRepos(SOURCE_ORG))
+                .ReturnsAsync(new[] { (REPO, "private") });
+
+            var expected = new StringBuilder();
+            expected.AppendLine($"$MigrationID = ExecAndGetMigrationID {{ gh gei migrate-repo --github-source-org \"{SOURCE_ORG}\" --source-repo \"{REPO}\" --github-target-org \"{TARGET_ORG}\" --target-repo \"{REPO}\" --ghes-api-url \"https://foo.com/api/v3\" --use-github-storage --queue-only --target-repo-visibility private }}");
+            expected.AppendLine($"$RepoMigrations[\"{REPO}\"] = $MigrationID");
+            expected.Append($"if ($RepoMigrations[\"{REPO}\"]) {{ gh gei wait-for-migration --migration-id $RepoMigrations[\"{REPO}\"] }}");
+
+            // Act
+            var args = new GenerateScriptCommandArgs
+            {
+                GithubSourceOrg = SOURCE_ORG,
+                GithubTargetOrg = TARGET_ORG,
+                Output = new FileInfo("unit-test-output"),
+                UseGithubStorage = true,
+                GhesApiUrl = "https://foo.com/api/v3",
+            };
+            await _handler.Handle(args);
+
+            _script = TrimNonExecutableLines(_script, 19, 7);
+
+            // Assert
+            _script.Should().Be(expected.ToString());
+        }
+
+
+
         [Fact]
         public async Task Validates_Env_Vars_Blob_Storage_Not_Validated_When_GHES_3_8()
         {

src/OctoshiftCLI.Tests/gei/Commands/GenerateScript/GenerateScriptCommandTests.cs

@@ -39,7 +39,7 @@ public void Should_Have_Options()
             var command = new GenerateScriptCommand();
             command.Should().NotBeNull();
             command.Name.Should().Be("generate-script");
-            command.Options.Count.Should().Be(15);
+            command.Options.Count.Should().Be(16);
 
             TestHelpers.VerifyCommandOption(command.Options, "github-source-org", true);
             TestHelpers.VerifyCommandOption(command.Options, "github-target-org", true);
@@ -56,6 +56,7 @@ public void Should_Have_Options()
             TestHelpers.VerifyCommandOption(command.Options, "aws-region", false);
             TestHelpers.VerifyCommandOption(command.Options, "keep-archive", false);
             TestHelpers.VerifyCommandOption(command.Options, "target-api-url", false);
+            TestHelpers.VerifyCommandOption(command.Options, "use-github-storage", false, true);
         }
 
         [Fact]

src/gei/Commands/GenerateScript/GenerateScriptCommand.cs

@@ -36,6 +36,7 @@ public GenerateScriptCommand() : base(
             AddOption(GithubSourcePat);
             AddOption(Verbose);
             AddOption(KeepArchive);
+            AddOption(UseGithubStorage);
         }
         public Option<string> GithubSourceOrg { get; } = new("--github-source-org")
         {
@@ -99,6 +100,11 @@ public GenerateScriptCommand() : base(
         {
             Description = "The URL of the target API, if not migrating to github.com. Defaults to https://api.github.com"
         };
+        public Option<bool> UseGithubStorage { get; } = new("--use-github-storage")
+        {
+            IsHidden = true,
+            Description = "Enables multipart uploads to a GitHub owned storage for use during migration",
+        };
 
         public override GenerateScriptCommandHandler BuildHandler(GenerateScriptCommandArgs args, IServiceProvider sp)
         {

src/gei/Commands/GenerateScript/GenerateScriptCommandArgs.cs

@@ -23,19 +23,33 @@ public class GenerateScriptCommandArgs : CommandArgs
         public string GithubSourcePat { get; set; }
         public bool KeepArchive { get; set; }
         public string TargetApiUrl { get; set; }
+        public bool UseGithubStorage { get; set; }
 
         public override void Validate(OctoLogger log)
         {
-            if (AwsBucketName.HasValue() && GhesApiUrl.IsNullOrWhiteSpace())
+            if (AwsBucketName.HasValue())
             {
-                throw new OctoshiftCliException("--ghes-api-url must be specified when --aws-bucket-name is specified.");
+                if (GhesApiUrl.IsNullOrWhiteSpace())
+                {
+                    throw new OctoshiftCliException("--ghes-api-url must be specified when --aws-bucket-name is specified.");
+                }
+
+                if (UseGithubStorage)
+                {
+                    throw new OctoshiftCliException("The --use-github-storage flag was provided with an AWS S3 Bucket name. Archive cannot be uploaded to both locations.");
+                }
             }
 
             if (NoSslVerify && GhesApiUrl.IsNullOrWhiteSpace())
             {
                 throw new OctoshiftCliException("--ghes-api-url must be specified when --no-ssl-verify is specified.");
             }
 
+            if (GhesApiUrl.IsNullOrWhiteSpace() && UseGithubStorage)
+            {
+                throw new OctoshiftCliException("--ghes-api-url must be specified when --use-github-storage is specified.");
+            }
+
             if (GhesApiUrl.HasValue())
             {
                 var result = Uri.TryCreate(GhesApiUrl, UriKind.Absolute, out var uriResult)