text: ##
  # Copyright (c) 2010-2017 Apple Inc. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  ##
  
  from twext.enterprise.dal.syntax import Delete, Insert, Select, Count
  
  from twisted.internet.defer import inlineCallbacks, succeed, returnValue
  from twisted.trial.unittest import TestCase
  
  from twistedcaldav.config import config
  
  from txdav.caldav.datastore.sql import CalendarStoreFeatures
  from txdav.caldav.datastore.test.util import CommonStoreTests
  from txdav.common.datastore.sql_tables import schema
  from txdav.common.datastore.test.util import theStoreBuilder, \
      StubNotifierFactory
  from txdav.common.datastore.upgrade.sql.others import attachment_migration
  from txdav.common.datastore.upgrade.sql.upgrade import UpgradeDatabaseOtherStep
  
  import hashlib
  import os
  
  """
  Tests for L{txdav.common.datastore.upgrade.sql.upgrade}.
  """
  
  
  class AttachmentMigrationModeTests(CommonStoreTests):
      """
      Tests for L{UpgradeDatabaseSchemaStep}.
      """
  
      @inlineCallbacks
      def setUp(self):
          # We need to skip the immediate base class since we are creating our own
          # store in each test
          yield TestCase.setUp(self)
  
      @inlineCallbacks
      def _initStore(self, enableManagedAttachments=True):
          """
          Build a store with certain bits cleaned out.
          """
  
          self.patch(config, "EnableManagedAttachments", enableManagedAttachments)
  
          store = yield theStoreBuilder.buildStore(
              self, {"push": StubNotifierFactory()}
          )
          store.enableManagedAttachments = enableManagedAttachments
  
          txn = store.newTransaction()
          cs = schema.CALENDARSERVER
          yield Delete(
              From=cs,
              Where=cs.NAME == "MANAGED-ATTACHMENTS"
          ).on(txn)
          yield txn.commit()
  
          returnValue(store)
  
      @inlineCallbacks
      def test_upgradeFromEmptyDropbox(self):
          """
          Test L{attachment_migration.doUpgrade} when managed attachments is enabled and dropbox items do not exist.
          """
  
          didUpgrade = [False, ]
  
          def _hasDropboxAttachments(_self, txn):
              return succeed(False)
          self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)
  
          def _upgradeToManagedAttachments(_self, batchSize=10):
              didUpgrade[0] = True
              return succeed(None)
          self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)
  
          store = (yield self._initStore())
  
          upgrader = UpgradeDatabaseOtherStep(store)
          yield attachment_migration.doUpgrade(upgrader)
          self.assertFalse(didUpgrade[0])
  
          txn = upgrader.sqlStore.newTransaction()
          managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
          yield txn.commit()
          self.assertNotEqual(managed, None)
  
      @inlineCallbacks
      def test_upgradeFromDropboxOK(self):
          """
          Test L{attachment_migration.doUpgrade} when managed attachments is enabled and dropbox items exist.
          """
  
          didUpgrade = [False, ]
  
          def _hasDropboxAttachments(_self, txn):
              return succeed(True)
          self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)
  
          def _upgradeToManagedAttachments(_self, batchSize=10):
              didUpgrade[0] = True
              return succeed(None)
          self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)
  
          store = (yield self._initStore())
  
          upgrader = UpgradeDatabaseOtherStep(store)
          yield attachment_migration.doUpgrade(upgrader)
          self.assertFalse(didUpgrade[0])
  
          txn = upgrader.sqlStore.newTransaction()
          managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
          yield txn.commit()
          self.assertNotEqual(managed, None)
  
      @inlineCallbacks
      def test_upgradeAlreadyDone(self):
          """
          Test L{attachment_migration.doUpgrade} when managed attachments is enabled and migration already done.
          """
  
          didUpgrade = [False, ]
  
          def _hasDropboxAttachments(_self, txn):
              return succeed(True)
          self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)
  
          def _upgradeToManagedAttachments(_self, batchSize=10):
              didUpgrade[0] = True
              return succeed(None)
          self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)
  
          store = (yield self._initStore())
          txn = store.newTransaction()
          yield txn.setCalendarserverValue("MANAGED-ATTACHMENTS", "1")
          yield txn.commit()
  
          upgrader = UpgradeDatabaseOtherStep(store)
          yield attachment_migration.doUpgrade(upgrader)
          self.assertFalse(didUpgrade[0])
  
          txn = upgrader.sqlStore.newTransaction()
          managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
          yield txn.commit()
          self.assertNotEqual(managed, None)
  
      @inlineCallbacks
      def test_upgradeNotEnabled(self):
          """
          Test L{attachment_migration.doUpgrade} when managed attachments is disabled.
          """
  
          didUpgrade = [False, ]
  
          def _hasDropboxAttachments(_self, txn):
              return succeed(True)
          self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)
  
          def _upgradeToManagedAttachments(_self, batchSize=10):
              didUpgrade[0] = True
              return succeed(None)
          self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)
  
          store = (yield self._initStore(False))
  
          upgrader = UpgradeDatabaseOtherStep(store)
          yield attachment_migration.doUpgrade(upgrader)
          self.assertFalse(didUpgrade[0])
  
          txn = upgrader.sqlStore.newTransaction()
          managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
          yield txn.commit()
          self.assertEqual(managed, None)
  
  
  class AttachmentMigrationTests(CommonStoreTests):
      """
      Tests for L{UpgradeDatabaseSchemaStep}.
      """
  
      @inlineCallbacks
      def setUp(self):
          self.patch(config, "EnableManagedAttachments", True)
  
          yield super(AttachmentMigrationTests, self).setUp()
  
          self._sqlCalendarStore.enableManagedAttachments = True
  
          txn = self.transactionUnderTest()
          cs = schema.CALENDARSERVER
          yield Delete(
              From=cs,
              Where=cs.NAME == "MANAGED-ATTACHMENTS"
          ).on(txn)
          yield self.commit()
  
      @inlineCallbacks
      def test_upgradeOrphanedAttachment(self):
          """
          Test L{attachment_migration.doUpgrade} when an orphaned attachment is present.
          """
  
          def _hasDropboxAttachments(_self, txn):
              return succeed(True)
          self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)
  
          # Create orphaned attachment
          dropboxID = "ABCD.dropbox"
          attachmentName = "test.txt"
          home = yield self.homeUnderTest(name="user01")
          at = schema.ATTACHMENT
          yield Insert(
              {
                  at.CALENDAR_HOME_RESOURCE_ID: home._resourceID,
                  at.DROPBOX_ID: dropboxID,
                  at.CONTENT_TYPE: "text/plain",
                  at.SIZE: 10,
                  at.MD5: "abcd",
                  at.PATH: attachmentName,
              }
          ).on(self.transactionUnderTest())
          yield self.commit()
  
          hasheduid = hashlib.md5(dropboxID).hexdigest()
          fp = self._sqlCalendarStore.attachmentsPath.child(hasheduid[0:2]).child(hasheduid[2:4]).child(hasheduid)
          fp.makedirs()
          fp = fp.child(attachmentName)
          fp.setContent("1234567890")
  
          self.assertTrue(os.path.exists(fp.path))
  
          upgrader = UpgradeDatabaseOtherStep(self._sqlCalendarStore)
          yield attachment_migration.doUpgrade(upgrader)
  
          txn = upgrader.sqlStore.newTransaction()
          managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
          count = (yield Select(
              [Count(at.DROPBOX_ID), ],
              From=at,
          ).on(txn))[0][0]
          yield txn.commit()
          self.assertEqual(count, 1)
          self.assertNotEqual(managed, None)
  
          self.assertTrue(os.path.exists(fp.path))
  