From 2f06cc6ffa7fdf90aeb692d56c19948df6319712 Mon Sep 17 00:00:00 2001
From: Taloth Saldono <Taloth@users.noreply.github.com>
Date: Sun, 25 Jan 2015 22:21:17 +0100
Subject: [PATCH] Fixed: Branch redirects will now occur during install of the
 a new update instead of during an update check.

---
 .../UpdateTests/UpdateServiceFixture.cs       | 12 +++++++++++
 .../Update/InstallUpdateService.cs            | 21 +++++++++++++++++++
 .../Update/UpdateCheckService.cs              | 15 -------------
 3 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/src/NzbDrone.Core.Test/UpdateTests/UpdateServiceFixture.cs b/src/NzbDrone.Core.Test/UpdateTests/UpdateServiceFixture.cs
index 29637d3cc..981fc894b 100644
--- a/src/NzbDrone.Core.Test/UpdateTests/UpdateServiceFixture.cs
+++ b/src/NzbDrone.Core.Test/UpdateTests/UpdateServiceFixture.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.IO;
 using FluentAssertions;
 using Moq;
@@ -294,6 +295,17 @@ namespace NzbDrone.Core.Test.UpdateTests
             ExceptionVerification.ExpectedErrors(1);
         }
 
+        [Test]
+        public void should_switch_to_branch_specified_in_updatepackage()
+        {
+            _updatePackage.Branch = "fake";
+
+            Subject.Execute(new InstallUpdateCommand() { UpdatePackage = _updatePackage });
+
+            Mocker.GetMock<IConfigFileProvider>()
+                  .Verify(v => v.SaveConfigDictionary(It.Is<Dictionary<string, object>>(d => d.ContainsKey("Branch") && (string)d["Branch"] == "fake")), Times.Once());
+        }
+
         [TearDown]
         public void TearDown()
         {
diff --git a/src/NzbDrone.Core/Update/InstallUpdateService.cs b/src/NzbDrone.Core/Update/InstallUpdateService.cs
index 7635fe88a..300eb6f48 100644
--- a/src/NzbDrone.Core/Update/InstallUpdateService.cs
+++ b/src/NzbDrone.Core/Update/InstallUpdateService.cs
@@ -102,6 +102,8 @@ namespace NzbDrone.Core.Update
             _archiveService.Extract(packageDestination, updateSandboxFolder);
             _logger.Info("Update package extracted successfully");
 
+            EnsureValidBranch(updatePackage);
+
             _backupService.Backup(BackupType.Update);
 
             if (OsInfo.IsNotWindows && _configFileProvider.UpdateMechanism == UpdateMechanism.Script)
@@ -120,6 +122,25 @@ namespace NzbDrone.Core.Update
             _processProvider.Start(_appFolderInfo.GetUpdateClientExePath(), GetUpdaterArgs(updateSandboxFolder));
         }
 
+        private void EnsureValidBranch(UpdatePackage package)
+        {
+            var currentBranch = _configFileProvider.Branch;
+            if (package.Branch != currentBranch)
+            {
+                try
+                {
+                    _logger.Info("Branch [{0}] is being redirected to [{1}]]", currentBranch, package.Branch);
+                    var config = new Dictionary<string, object>();
+                    config["Branch"] = package.Branch;
+                    _configFileProvider.SaveConfigDictionary(config);
+                }
+                catch (Exception e)
+                {
+                    _logger.ErrorException(string.Format("Couldn't change the branch from [{0}] to [{1}].", currentBranch, package.Branch), e);
+                }
+            }
+        }
+
         private void InstallUpdateWithScript(String updateSandboxFolder)
         {
             var scriptPath = _configFileProvider.UpdateScriptPath;
diff --git a/src/NzbDrone.Core/Update/UpdateCheckService.cs b/src/NzbDrone.Core/Update/UpdateCheckService.cs
index 16ffe17da..f17b72dcc 100644
--- a/src/NzbDrone.Core/Update/UpdateCheckService.cs
+++ b/src/NzbDrone.Core/Update/UpdateCheckService.cs
@@ -42,21 +42,6 @@ namespace NzbDrone.Core.Update
             {
                 _logger.ProgressDebug("No update available.");
             }
-            else if (latestAvailable.Branch != _configFileProvider.Branch)
-            {
-                try
-                {
-                    _logger.Info("Branch [{0}] is being redirected to [{1}]]", _configFileProvider.Branch, latestAvailable.Branch);
-                    var config = new Dictionary<string, object>();
-                    config["Branch"] = latestAvailable.Branch;
-                    _configFileProvider.SaveConfigDictionary(config);
-                }
-                catch (Exception e)
-                {
-                    _logger.ErrorException("Couldn't save the branch redirect.", e);
-                }
-
-            }
 
             return latestAvailable;
         }