package com.atlassian.sal.core.upgrade;
   
   import com.atlassian.beehive.ClusterLock;
   import com.atlassian.beehive.ClusterLockService;
   import com.atlassian.plugin.Plugin;
   import com.atlassian.plugin.PluginAccessor;
   import com.atlassian.plugin.event.PluginEventManager;
   import com.atlassian.sal.api.message.Message;
   import com.atlassian.sal.api.pluginsettings.PluginSettings;
  import com.atlassian.sal.api.pluginsettings.PluginSettingsFactory;
  import com.atlassian.sal.api.transaction.TransactionCallback;
  import com.atlassian.sal.api.transaction.TransactionTemplate;
  import com.atlassian.sal.api.upgrade.PluginUpgradeTask;
  import com.atlassian.sal.core.message.DefaultMessage;
  
  import org.junit.After;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.mockito.ArgumentCaptor;
  import org.mockito.Mock;
  import org.mockito.runners.MockitoJUnitRunner;
  
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.List;
  import java.util.concurrent.TimeUnit;
  
  import static com.atlassian.sal.core.upgrade.PluginUpgrader.BUILD;
  import static java.util.Collections.singletonList;
  import static org.hamcrest.Matchers.contains;
  import static org.hamcrest.Matchers.empty;
  import static org.hamcrest.Matchers.is;
  import static org.hamcrest.Matchers.not;
  import static org.junit.Assert.assertThat;
  import static org.mockito.Matchers.anyString;
  import static org.mockito.Mockito.never;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  
  @RunWith (MockitoJUnitRunner.class)
  public class DefaultPluginUpgradeManagerTest
  {
      @Mock private Plugin plugin;
      @Mock private Plugin secondPlugin;
      @Mock private PluginUpgradeTask pluginUpgradeTask;
      @Mock private PluginUpgradeTask secondPluginUpgradeTask;
      @Mock private PluginUpgradeTask thirdPluginUpgradeTask;
      @Mock private PluginAccessor pluginAccessor;
      @Mock private PluginSettingsFactory pluginSettingsFactory;
      @Mock private PluginEventManager pluginEventManager;
      @Mock private PluginSettings pluginSettings;
      @Mock private ClusterLockService clusterLockService;
      @Mock private ClusterLock clusterLock;
  
      private static final String PLUGIN_KEY = "com.atlassian.plugin.test.upgrade-test";
      private static final String SECOND_PLUGIN_KEY = "com.atlassian.plugin2.test.upgrade-test";
  
      private DefaultPluginUpgradeManager pluginUpgradeManager;
      private CountingTransactionTemplate transactionTemplate;
  
      @Before
      public void setUp() throws Exception
      {
          transactionTemplate = new CountingTransactionTemplate();
  
          when(plugin.getKey()).thenReturn(PLUGIN_KEY);
          when(secondPlugin.getKey()).thenReturn(SECOND_PLUGIN_KEY);
  
          when(pluginAccessor.getPlugin(PLUGIN_KEY)).thenReturn(plugin);
          when(pluginAccessor.getPlugin(SECOND_PLUGIN_KEY)).thenReturn(secondPlugin);
  
          when(pluginUpgradeTask.getPluginKey()).thenReturn(PLUGIN_KEY);
          when(pluginUpgradeTask.getBuildNumber()).thenReturn(1);
          when(pluginUpgradeTask.doUpgrade()).thenReturn(Collections.<Message>emptyList()); // default is success
  
          when(secondPluginUpgradeTask.getPluginKey()).thenReturn(PLUGIN_KEY);
          when(secondPluginUpgradeTask.getBuildNumber()).thenReturn(2);
          when(secondPluginUpgradeTask.doUpgrade()).thenReturn(Collections.<Message>emptyList()); // default is success
  
          when(thirdPluginUpgradeTask.getPluginKey()).thenReturn(SECOND_PLUGIN_KEY);
          when(thirdPluginUpgradeTask.getBuildNumber()).thenReturn(1);
          when(thirdPluginUpgradeTask.doUpgrade()).thenReturn(Collections.<Message>emptyList()); // default is success
  
          when(pluginSettingsFactory.createGlobalSettings()).thenReturn(pluginSettings);
  
          when(clusterLockService.getLockForName(anyString())).thenReturn(clusterLock);
      }
  
      @After
      public void tearDown()
      {
          transactionTemplate.resetTransactionCount();
      }
  
      @Test
      public void testUpgradeInternalIsSuccessful() throws Exception
      {
         pluginUpgradeManager = new DefaultPluginUpgradeManager(singletonList(pluginUpgradeTask), transactionTemplate,
                 pluginAccessor, pluginSettingsFactory, pluginEventManager, clusterLockService);
 
         when(pluginSettings.get(plugin.getKey() + BUILD)).thenReturn(Integer.toString(0)); // last upgrade task (data build number)
         when(clusterLock.tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS)).thenReturn(true);
 
         List<Message> messages = pluginUpgradeManager.upgradeInternal(plugin);
 
         // A successful upgrade task should increment the build data number by a put on plugin settings
         verify(pluginSettings).put(PLUGIN_KEY + BUILD, "1");
 
         assertThat(messages, empty());
 
         verify(clusterLock).tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS);
         verify(clusterLock).unlock();
     }
 
     @Test
     public void testUpgradeInternalFailsIfUpgradeFails() throws Exception
     {
         pluginUpgradeManager = new DefaultPluginUpgradeManager(singletonList(pluginUpgradeTask), transactionTemplate,
                 pluginAccessor, pluginSettingsFactory, pluginEventManager, clusterLockService);
 
         when(pluginUpgradeTask.doUpgrade()).thenReturn(Collections.<Message>singletonList(new DefaultMessage("failed.upgrade")));
         when(clusterLock.tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS)).thenReturn(true);
 
         List<Message> messages = pluginUpgradeManager.upgradeInternal(plugin);
 
         // A failed upgrade task should not increment the build data number in plugin settings
         verify(pluginSettings, never()).put(anyString(), anyString());
 
         assertThat(messages, not(empty()));
 
         ArgumentCaptor<String> lockName = ArgumentCaptor.forClass(String.class);
         verify(clusterLockService).getLockForName(lockName.capture());
         verify(clusterLock).tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS);
         verify(clusterLock).unlock();
 
         messages = pluginUpgradeManager.upgradeInternal(plugin);
 
         verify(clusterLockService, times(2)).getLockForName(lockName.getValue());
     }
 
     @Test
     public void testThatOnlyOneTransactionTemplateIsExecutedPerPluginUpgrade() throws Exception
     {
         pluginUpgradeManager = new DefaultPluginUpgradeManager(Arrays.asList(pluginUpgradeTask, secondPluginUpgradeTask), transactionTemplate,
                 pluginAccessor, pluginSettingsFactory, pluginEventManager, clusterLockService);
 
         when(clusterLock.tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS)).thenReturn(true);
 
         List<Message> messages = pluginUpgradeManager.upgradeInternal(plugin);
 
         assertThat(messages, empty());
         assertThat(transactionTemplate.getTransactionCount(), is(1));
 
         verify(pluginSettings).put(PLUGIN_KEY + BUILD, "2"); // from pluginUpgradeTask, secondPluginUpgradeTask
 
         verify(clusterLock).tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS);
         verify(clusterLock).unlock();
     }
 
     @Test
     public void testThatOnlyOneTransactionTemplateIsExecutedPerPluginUpgradeWithMultiplePlugins() throws Exception
     {
         pluginUpgradeManager = new DefaultPluginUpgradeManager(Arrays.asList(pluginUpgradeTask, secondPluginUpgradeTask, thirdPluginUpgradeTask), transactionTemplate,
                 pluginAccessor, pluginSettingsFactory, pluginEventManager, clusterLockService);
 
         when(clusterLock.tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS)).thenReturn(true);
 
         List<Message> messages = pluginUpgradeManager.upgradeInternal();
 
         verify(pluginSettings).put(PLUGIN_KEY + BUILD, "2"); // from pluginUpgradeTask, secondPluginUpgradeTask
         verify(pluginSettings).put(SECOND_PLUGIN_KEY + BUILD, "1"); // thirdPluginUpgradeTask
 
         assertThat(messages, empty());
         assertThat(transactionTemplate.getTransactionCount(), is(2));
 
         verify(clusterLock, times(2)).tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS);
         verify(clusterLock, times(2)).unlock();
     }
 
     @Test
     public void testUpgradeLockTimeout() throws InterruptedException
     {
         pluginUpgradeManager = new DefaultPluginUpgradeManager(singletonList(pluginUpgradeTask), transactionTemplate,
                 pluginAccessor, pluginSettingsFactory, pluginEventManager, clusterLockService);
 
         when(clusterLock.tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS)).thenReturn(false);
 
         List<Message> messages = pluginUpgradeManager.upgradeInternal(plugin);
 
         final Message expectedMessage = new DefaultMessage("unable to acquire cluster lock named 'sal.upgrade.com.atlassian.plugin.test.upgrade-test' after waiting " + DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS + " seconds; note that this timeout may be adjusted via the system property '" + DefaultPluginUpgradeManager.LOCK_TIMEOUT_PROPERTY + "'");
         assertThat(messages, contains(expectedMessage));
 
         verify(clusterLock).tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS);
     }
 
     @Test
     public void testUpgradeLockInterrupted() throws InterruptedException
     {
         pluginUpgradeManager = new DefaultPluginUpgradeManager(singletonList(pluginUpgradeTask), transactionTemplate,
                 pluginAccessor, pluginSettingsFactory, pluginEventManager, clusterLockService);
 
         when(clusterLock.tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS)).thenThrow(new InterruptedException("blargh"));
 
         List<Message> messages = pluginUpgradeManager.upgradeInternal(plugin);
 
         final Message expectedMessage = new DefaultMessage("interrupted while trying to acquire cluster lock named 'sal.upgrade.com.atlassian.plugin.test.upgrade-test' blargh");
         assertThat(messages, contains(expectedMessage));
 
         verify(clusterLock).tryLock(DefaultPluginUpgradeManager.LOCK_TIMEOUT_SECONDS, TimeUnit.SECONDS);
     }
 
     private static class CountingTransactionTemplate implements TransactionTemplate
     {
         private int transactionCount = 0;
 
         @Override
         public <T> T execute(final TransactionCallback<T> action)
         {
             transactionCount++;
             return action.doInTransaction();
         }
 
         void resetTransactionCount()
         {
             transactionCount = 0;
         }
 
         int getTransactionCount()
         {
             return transactionCount;
         }
     }
 }