content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE TypeApplications #-}
  module CatalogTest (main) where
  
  import Control.Concurrent.MVar
  import Util.STM
  import Control.Exception hiding (assert)
  import Control.Monad
  import Data.Default
  import qualified Data.HashMap.Strict as HashMap
  import Data.HashSet (HashSet)
  import qualified Data.HashSet as HashSet
  import Data.List
  import Data.Proxy
  import qualified Data.Text as Text
  import GHC.Stack (HasCallStack)
  import Test.HUnit
  import Test.QuickCheck
  import Test.QuickCheck.IO ()
  
  import TestRunner
  import Util.Testing
  
  import Glean.Database.Catalog (Catalog)
  import qualified Glean.Database.Catalog as Catalog
  import Glean.Database.Catalog.Filter
  import qualified Glean.Database.Catalog.Store as Store
  import Glean.Database.Catalog.Test
  import Glean.Init
  import Glean.Test.HUnit
  import Glean.Repo.Text
  import Glean.Test.Mock
  import Glean.Internal.Types
  import Glean.Types hiding (Exception)
  
  repos :: [Repo]
  repos = [Repo "foo" (Text.pack $ show n) | n <- [1 .. 12 :: Int]]
  
  data CreateDelete = Create | Delete
    deriving(Eq,Ord,Enum,Bounded,Read,Show)
  
  instance Arbitrary CreateDelete where
    arbitrary = elements [Create,Delete]
  
  newtype ARepo = ARepo Repo
    deriving(Eq)
  
  instance Show ARepo where
    show (ARepo repo) = showRepo repo
  
  instance Arbitrary ARepo where
    arbitrary = ARepo <$> elements repos
  
  data CreateDeleteAction = CreateDeleteAction
    { cdaWhat :: CreateDelete
    , cdaRepo :: Repo
    , cdaExists :: Bool
    , cdaAfter :: HashSet Repo
    }
  
  populate :: [(CreateDelete, ARepo)] -> [CreateDeleteAction]
  populate = snd . mapAccumL mk mempty
    where
      mk before (what, ARepo repo) =
        let after = (if what == Create then HashSet.insert else HashSet.delete)
              repo
              before
        in
        ( after
        , CreateDeleteAction
            { cdaWhat = what
            , cdaRepo = repo
            , cdaExists = repo `HashSet.member` before
            , cdaAfter = after
            } )
  
  createDeleteTest :: Test
  createDeleteTest = TestCase $ assertProperty "" $ test . populate
    where
      test trace = withMemCatalog $ \_ cat ->
        forM_ trace $ \CreateDeleteAction{..} -> do
          exec cdaWhat cdaExists cat cdaRepo
          xs <- atomically $ Catalog.list cat [Local] everythingF
          assertEqual "" cdaAfter $ HashSet.fromList (map itemRepo xs)
          checkConsistency cat
  
      exec Create False = createNew
      exec Create True = createExisting
      exec Delete False = deleteUnknown
      exec Delete True = deleteExisting
  
  data TestError = TestError
    deriving(Eq, Show)
  
  instance Exception TestError
  
  createFailureTest :: Test
  createFailureTest = TestCase $ withMemCatalog $ \store cat -> do
    prepare (storeCreate store) [ doBefore $ throwIO TestError ]
    called <- newTVarIO False
    assertThrows "" TestError $
      Catalog.create cat repo1 def $ writeTVar called True
    assert $ not <$> readTVarIO called
    assert $ not <$> Catalog.exists cat [Local] repo1
    checkConsistency cat
    where
      repo1 = head repos
  
  nextMeta :: Meta -> Meta
  nextMeta !meta = meta
    { metaProperties = HashMap.insert
        "counter"
        (Text.pack $ show $ metaCounter meta + 1)
        (metaProperties meta)
    }
  
  metaCounter :: Meta -> Int
  metaCounter =
    maybe 0 (read . Text.unpack) . HashMap.lookup "counter" . metaProperties
  
  writeMetaTest :: Test
  writeMetaTest = TestCase $ withMemCatalog $ \store cat -> do
    Catalog.create cat repo1 def $ return ()
  
    let writes = [5,1,64,1024]
  
    val <- newTVarIO 0
    max <- newTVarIO 0
  
    augment (storePut store) $ \impl repo meta -> do
      let n = metaCounter meta
      m <- readTVarIO max
      assert $ n <= m
      b <- impl repo meta
      atomically $ writeTVar val n
      return b
  
    forM_ writes $ \n -> when (n /= 0) $ do
      k <- atomically $ metaCounter <$> Catalog.readMeta cat repo1
      atomically $ writeTVar max $ k+n
      replicateM_ n $ void $ atomically $
          metaCounter <$> Catalog.modifyMeta cat repo1 (return . nextMeta)
      atomically $ do
        m <- readTVar val
        when (m /= k+n) retry
    where
      repo1 = head repos
  
  stackedDbsTest :: Test
  stackedDbsTest = TestCase $ withMemCatalog $ \_ cat -> do
    Catalog.create cat repoA metaA $ return ()
    Catalog.create cat repoC metaC $ return ()
  
    assertEqual "" DatabaseStatus_Complete =<< status cat repoA
    assertThrows "" TestError $ status cat repoB
    assertEqual "" DatabaseStatus_Missing =<< status cat repoC
  
    Catalog.create cat repoB metaB $ return ()
  
    assertEqual "" DatabaseStatus_Complete =<< status cat repoA
    assertEqual "" DatabaseStatus_Complete =<< status cat repoB
    assertEqual "" DatabaseStatus_Complete =<< status cat repoC
  
    Catalog.delete cat repoB
  
    assertEqual "" DatabaseStatus_Complete =<< status cat repoA
    assertThrows "" TestError $ status cat repoB
    assertEqual "" DatabaseStatus_Missing =<< status cat repoC
  
    atomically $ Catalog.startRestoring cat repoB metaB
  
    assertEqual "" DatabaseStatus_Complete =<< status cat repoA
    assertEqual "" DatabaseStatus_Restoring =<< status cat repoB
    assertEqual "" DatabaseStatus_Restoring =<< status cat repoC
  
    Catalog.finishRestoring cat repoB
  
    assertEqual "" DatabaseStatus_Complete =<< status cat repoA
    assertEqual "" DatabaseStatus_Complete =<< status cat repoB
    assertEqual "" DatabaseStatus_Complete =<< status cat repoC
  
    Catalog.delete cat repoC
    Catalog.create cat repoC metaCWrongGuid $ return ()
  
    assertEqual "" DatabaseStatus_Complete =<< status cat repoA
    assertEqual "" DatabaseStatus_Complete =<< status cat repoB
    assertEqual "" DatabaseStatus_Missing =<< status cat repoC
  
    Catalog.delete cat repoC
    Catalog.create cat repoC metaC $ return ()
  
    assertEqual "" DatabaseStatus_Complete =<< status cat repoA
    assertEqual "" DatabaseStatus_Complete =<< status cat repoB
    assertEqual "" DatabaseStatus_Complete =<< status cat repoC
  
    where
      -- Dependencies: C -> B -> A
      repoA = Repo "base" "repoA"
      repoB = Repo "stacked" "repoB"
      repoC = Repo "stacked2" "repoC"
      guidA = "aaaaaa"
      guidB = "bbbbbb"
      guidC = "cccccc"
      guidWrong = "xxxxxx"
      stacked (Repo name hash) guid =
        Dependencies_stacked $ Stacked name hash (Just guid)
      metaA = meta
        { metaProperties=HashMap.fromList [("glean.guid", guidA)] }
      metaB = meta
        { metaDependencies=Just $ stacked repoA guidA
        , metaProperties=HashMap.fromList [("glean.guid", guidB)]}
      metaC = meta
        { metaDependencies=Just $ stacked repoB guidB
        , metaProperties=HashMap.fromList [("glean.guid", guidC)]}
      metaCWrongGuid = meta
        { metaDependencies=Just $ stacked repoB guidWrong
        , metaProperties=HashMap.fromList [("glean.guid", guidC)]}
      meta = def
        { metaCompleteness=Complete $ DatabaseComplete (PosixEpochTime 0) Nothing}
  
      status cat repo = do
        result <- atomically $ Catalog.getLocalDatabase cat repo
        case result of
          Just GetDatabaseResult{getDatabaseResult_database=db} ->
            return (database_status db)
          Nothing -> throwIO TestError
  
  writeWhileCommitting :: Test
  writeWhileCommitting = TestCase $ withMemCatalog $ \store cat -> do
    Catalog.create cat repo1 def $ return ()
  
    start1 <- newEmptyMVar
    finish1 <- newEmptyMVar
    finish2 <- newEmptyMVar
    prepare (storePut store)
      [ doBefore (putMVar start1 ()) . doFinally (takeMVar finish1)
      , doFinally (putMVar finish2 ()) ]
    void $ atomically $ Catalog.modifyMeta cat repo1 (return . nextMeta)
    takeMVar start1
    void $ atomically $ Catalog.modifyMeta cat repo1 (return . nextMeta)
    putMVar finish1 ()
    takeMVar finish2
    Just m <- Store.get store repo1
    assertEqual "" 2 $ metaCounter m
    where
      repo1 = head repos
  
  createNew :: HasCallStack => Catalog -> Repo -> IO ()
  createNew cat repo = do
    called <- newTVarIO False
    Catalog.create cat repo def $ writeTVar called True
    assert $ readTVarIO called
  
  createExisting :: HasCallStack => Catalog -> Repo -> IO ()
  createExisting cat repo = do
    called <- newTVarIO False
    assertThrows "" (Catalog.EntryAlreadyExists repo)
      $ Catalog.create cat repo def $ writeTVar called True
    assert $ not <$> readTVarIO called
  
  deleteExisting :: HasCallStack => Catalog -> Repo -> IO ()
  deleteExisting cat repo = Catalog.delete cat repo
  
  deleteUnknown :: HasCallStack => Catalog -> Repo -> IO ()
  deleteUnknown cat repo = assertThrowsType "" (Proxy @UnknownDatabase)
    $ Catalog.delete cat repo
  
  main :: IO ()
  main = withUnitTest $ testRunner $ TestList
    [ TestLabel "createDelete" createDeleteTest
    , TestLabel "createFailureTest" createFailureTest
    , TestLabel "writeMeta" writeMetaTest
    , TestLabel "stackedDbsTest" stackedDbsTest
    , TestLabel "writeWhileCommitting" writeWhileCommitting
    ]
  