ReindexSpec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Test.ReindexSpec (spec) where

import Control.Concurrent
import Numeric.Natural
import TestsUtils.Common
import TestsUtils.Import

spec :: Spec
spec =
  describe "Reindex" $ do
    it "reindexes synchronously" $ withTestEnv $ do
      _ <- insertData
      let newIndex = [qqIndexName|bloodhound-tests-twitter-index-1-reindexed|]
      _ <- tryPerformBHRequest $ deleteIndex newIndex
      _ <- performBHRequest $ createIndex (IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits) newIndex
      _ <- assertRight =<< tryPerformBHRequest (reindex $ mkReindexRequest testIndex newIndex)
      _ <- performBHRequest $ refreshIndex newIndex
      searchResult <- assertRight =<< tryPerformBHRequest (searchByIndex newIndex $ mkSearch Nothing Nothing)
      liftIO $
        map hitSource (hits (searchHits searchResult)) `shouldBe` [Just exampleTweet]

    it "reindexes asynchronously" $ withTestEnv $ do
      _ <- insertData
      let newIndex = [qqIndexName|bloodhound-tests-twitter-index-1-reindexed|]
      _ <- tryPerformBHRequest $ deleteIndex newIndex
      _ <- performBHRequest $ createIndex (IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits) newIndex
      taskNodeId <- assertRight =<< tryPerformBHRequest (reindexAsync $ mkReindexRequest testIndex newIndex)
      _ <- waitForTaskToComplete 10 taskNodeId
      _ <- performBHRequest $ refreshIndex newIndex
      searchResult <- assertRight =<< tryPerformBHRequest (searchByIndex newIndex (mkSearch Nothing Nothing))
      liftIO $
        (map hitSource (hits (searchHits searchResult))) `shouldBe` [Just exampleTweet]

    it "reindexWith forwards slices/max_docs URI params and reindexes synchronously" $ withTestEnv $ do
      -- Seed the source with two distinct documents so rioMaxDocs has
      -- something to cap. (insertData resets the index each call, so
      -- use insertTweetWithDocId for the second insertion.)
      _ <- insertData
      _ <- insertTweetWithDocId (exampleTweetWithAge 99) "2"
      let newIndex = [qqIndexName|bloodhound-tests-twitter-index-1-reindexed|]
          opts = defaultReindexOptions {rioSlices = Just SlicesAuto, rioMaxDocs = Just 1}
      _ <- tryPerformBHRequest $ deleteIndex newIndex
      _ <- performBHRequest $ createIndex (IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits) newIndex
      _ <- assertRight =<< tryPerformBHRequest (reindexWith opts $ mkReindexRequest testIndex newIndex)
      _ <- performBHRequest $ refreshIndex newIndex
      searchResult <- assertRight =<< tryPerformBHRequest (searchByIndex @Tweet newIndex (mkSearch Nothing Nothing))
      -- rioMaxDocs = 1 caps the destination to one document even though
      -- the source has two.
      liftIO $ length (hits (searchHits searchResult)) `shouldBe` 1

    it "rethrottleReindex unthrottles an in-progress async reindex" $ withTestEnv $ do
      _ <- insertData
      let newIndex = [qqIndexName|bloodhound-tests-twitter-rethrottle-target|]
      _ <- tryPerformBHRequest $ deleteIndex newIndex
      _ <-
        performBHRequest $
          createIndex (IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits) newIndex
      -- Start the reindex throttled to 1 doc/sec so the task is still
      -- running when the rethrottle lands.
      let opts = defaultReindexOptions {rioRequestsPerSecond = Just 1}
      taskNodeId <- assertRight =<< tryPerformBHRequest (reindexAsyncWith opts $ mkReindexRequest testIndex newIndex)
      -- Race window: the single-doc reindex may finish before the
      -- rethrottle reaches the server, in which case the response
      -- carries an empty node list (the docstring on 'rethrottleReindex'
      -- spells this out). When the task is still running the rethrottle
      -- response lists it (>= 1 node). Either outcome is acceptable
      -- here — the test exercises the wire path and decoder regardless.
      _ <- performBHRequest $ rethrottleReindex taskNodeId RethrottleUnlimited
      _ <- waitForTaskToComplete 10 taskNodeId
      _ <- performBHRequest $ refreshIndex newIndex
      searchResult <- assertRight =<< tryPerformBHRequest (searchByIndex newIndex (mkSearch Nothing Nothing))
      liftIO $
        map hitSource (hits (searchHits searchResult)) `shouldBe` [Just exampleTweet]

assertRight :: (Show a, MonadFail m) => Either a b -> m b
assertRight (Left x) = fail $ "Expected Right, got Left: " <> show x
assertRight (Right x) = pure x

-- | The response is not used, but make sure we return it so it gets parsed
waitForTaskToComplete :: (MonadBH m, MonadFail m) => Natural -> TaskNodeId -> m (TaskResponse ReindexResponse)
waitForTaskToComplete 0 taskNodeId = fail $ "Timed out waiting for task to complete, taskNodeId = " <> show taskNodeId
waitForTaskToComplete n taskNodeId = do
  task <- assertRight =<< tryPerformBHRequest (getTask taskNodeId)
  if taskResponseCompleted task
    then return task
    else liftIO (threadDelay 100000) >> waitForTaskToComplete (n - 1) taskNodeId