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.
  -}
  
  module PublishShardsTest (main) where
  
  import Control.Concurrent.Async
  import Util.STM
  import Data.Foldable (toList)
  import Data.Time
  import DatabaseJanitorTest (
    dbConfig,
    makeFakeDB,
    serverConfig,
    withTest,
   )
  import Glean.Database.Env
  import Glean.Database.Schema.Types (
    DbSchema,
   )
  import Glean.Init (withUnitTest)
  import Glean.Internal.Types
  import Glean.Server.Sharding
  import Glean.Types (Repo (Repo), DatabaseComplete (DatabaseComplete))
  import Test.HUnit
  import TestRunner
  import Glean.Database.Meta (utcTimeToPosixEpochTime)
  
  main :: IO ()
  main =
    withUnitTest $
      testRunner $
        TestList
          [ TestLabel "publish incomplete DBs" publishIncompleteDBs
          , TestLabel "publish complete DBs" publishCompleteDBs
          , TestLabel "stop publishing after timeout" stopPublishingCompletedDBs
          ]
  
  runTest :: Bool -> (IO [Repo] -> IO ()) -> IO ()
  runTest term test = do
    published <- newTVarIO mempty
    let callback = atomically . writeTVar published
    terminating <- newTVarIO term
    withTest setupFakeDBs setupCloudDBs $ \evb cfgAPI dbdir backupdir ->
      withDatabases evb (dbConfig dbdir $ serverConfig backupdir) cfgAPI $ \env ->
        withAsync (
            dbUpdateNotifierThread
            env
            0.0000001
            (readTVar terminating)
            callback
          ) $ \_ ->
              test $
                atomically $ do
                  repos <- readTVar published
                  if null repos
                    then retry
                    else return $ toList repos
  
  setupCloudDBs :: FilePath -> DbSchema -> IO ()
  setupCloudDBs _ _ = return ()
  
  setupFakeDBs :: FilePath -> DbSchema -> IO ()
  setupFakeDBs dbdir schema = do
    now <- getCurrentTime
    -- populate a dir with various DBs
    makeFakeDB schema dbdir incompleteRepo now incomplete id
    makeFakeDB schema dbdir completeRepo now complete id
  
  incomplete :: p -> Completeness
  incomplete _ = Incomplete DatabaseIncomplete
  
  complete :: UTCTime -> Completeness
  complete = Complete . (`DatabaseComplete` Nothing) . utcTimeToPosixEpochTime
  
  incompleteRepo :: Repo
  incompleteRepo = Repo "test" "incomplete"
  
  completeRepo :: Repo
  completeRepo = Repo "test" "complete"
  
  publishIncompleteDBs :: Test
  publishIncompleteDBs = TestCase $
    runTest False $ \readShards -> do
      repos <- readShards
      assertBool "Incomplete DBs are published" (incompleteRepo `elem` repos)
  
  publishCompleteDBs :: Test
  publishCompleteDBs = TestCase $
    runTest False $ \readShards -> do
      repos <- readShards
      assertBool "Complete DBs are published" (completeRepo `elem` repos)
  
  stopPublishingCompletedDBs :: Test
  stopPublishingCompletedDBs = TestCase $
    runTest True $ \readShards -> do
      repos <- readShards
      assertBool
        "Still publishing incomplete DBs after timeout"
        (incompleteRepo `elem` repos)
      assertBool
        "Stopped publishing completed DBs after timeout"
        (completeRepo `notElem` repos)
  