Common.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}

module TestsUtils.Common where

import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Versions as Versions
import Lens.Micro (toListOf)
import qualified Network.HTTP.Types.Status as NHTS
import TestsUtils.Import

testServer :: Server
testServer = Server "http://localhost:9200"

testIndex :: IndexName
testIndex = [qqIndexName|bloodhound-tests-twitter-1|]

withTestEnv :: BH IO a -> IO a
withTestEnv = withBH defaultManagerSettings testServer

data Location = Location
  { lat :: Double,
    lon :: Double
  }
  deriving stock (Eq, Show)

data Tweet = Tweet
  { user :: Text,
    postDate :: UTCTime,
    message :: Text,
    age :: Int,
    location :: Location,
    extra :: Maybe Text
  }
  deriving stock (Eq, Show)

$(deriveJSON defaultOptions ''Location)
$(deriveJSON defaultOptions ''Tweet)

data ConversationMapping = ConversationMapping deriving stock (Eq, Show)

instance ToJSON ConversationMapping where
  toJSON ConversationMapping =
    object
      [ "properties"
          .= object
            [ "reply_join"
                .= object
                  [ "type" .= ("join" :: Text),
                    "relations" .= object ["message" .= ("reply" :: Text)]
                  ],
              "user"
                .= object
                  [ "type" .= ("text" :: Text),
                    "fielddata" .= True
                  ],
              -- Serializing the date as a date is breaking other tests, mysteriously.
              -- , "postDate" .= object [ "type"   .= ("date" :: Text)
              --                        , "format" .= ("YYYY-MM-dd`T`HH:mm:ss.SSSZZ" :: Text)]
              "message" .= object ["type" .= ("text" :: Text)],
              "age" .= object ["type" .= ("integer" :: Text)],
              "location" .= object ["type" .= ("geo_point" :: Text)],
              "extra" .= object ["type" .= ("keyword" :: Text)]
            ]
      ]

getServerVersion :: IO Versions.Version
getServerVersion = extractVersion <$> withTestEnv (performBHRequest getStatus)
  where
    extractVersion = versionNumber . number . version

createExampleIndex :: (MonadBH m) => m (BHResponse StatusDependant Acknowledged, Acknowledged)
createExampleIndex = do
  result <- tryPerformBHRequest (keepBHResponse $ createIndex (IndexSettings (ShardCount 1) (ReplicaCount 0) defaultIndexMappingsLimits) testIndex)
  case result of
    Left e
      | T.isSuffixOf "already exists" (errorMessage e) -> return (error "TODO rewrite this part too", Acknowledged False)
      | otherwise -> throwEsError e
    Right ack -> return ack

deleteExampleIndex :: (MonadBH m) => m (BHResponse StatusDependant Acknowledged, Acknowledged)
deleteExampleIndex =
  performBHRequest $ keepBHResponse $ deleteIndex testIndex

validateStatus :: (Show body) => BHResponse contextualized body -> Int -> Expectation
validateStatus resp expected =
  if actual == expected
    then return ()
    else expectationFailure ("Expected " <> show expected <> " but got " <> show actual <> ": " <> show body)
  where
    actual = NHTS.statusCode (responseStatus $ getResponse resp)
    body = responseBody $ getResponse resp

data TweetMapping = TweetMapping deriving stock (Eq, Show)

instance ToJSON TweetMapping where
  toJSON TweetMapping =
    object
      [ "properties"
          .= object
            [ "user"
                .= object
                  [ "type" .= ("text" :: Text),
                    "fielddata" .= True
                  ],
              -- Serializing the date as a date is breaking other tests, mysteriously.
              -- , "postDate" .= object [ "type"   .= ("date" :: Text)
              --                        , "format" .= ("YYYY-MM-dd`T`HH:mm:ss.SSSZZ" :: Text)]
              "message" .= object ["type" .= ("text" :: Text)],
              "age" .= object ["type" .= ("integer" :: Text)],
              "location" .= object ["type" .= ("geo_point" :: Text)],
              "extra" .= object ["type" .= ("keyword" :: Text)]
            ]
      ]

exampleTweet :: Tweet
exampleTweet =
  Tweet
    { user = "bitemyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 10),
      message = "Use haskell!",
      age = 10000,
      location = Location 40.12 (-71.34),
      extra = Nothing
    }

tweetWithExtra :: Tweet
tweetWithExtra =
  Tweet
    { user = "bitemyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 10),
      message = "Use haskell!",
      age = 10000,
      location = Location 40.12 (-71.34),
      extra = Just "blah blah"
    }

exampleTweetWithAge :: Int -> Tweet
exampleTweetWithAge age' =
  Tweet
    { user = "bitemyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 10),
      message = "Use haskell!",
      age = age',
      location = Location 40.12 (-71.34),
      extra = Nothing
    }

newAge :: Int
newAge = 31337

newUser :: Text
newUser = "someotherapp"

tweetPatch :: Value
tweetPatch =
  object
    [ "age" .= newAge,
      "user" .= newUser
    ]

patchedTweet :: Tweet
patchedTweet = exampleTweet {age = newAge, user = newUser}

otherTweet :: Tweet
otherTweet =
  Tweet
    { user = "notmyapp",
      postDate =
        UTCTime
          (ModifiedJulianDay 55000)
          (secondsToDiffTime 11),
      message = "Use haskell!",
      age = 1000,
      location = Location 40.12 (-71.34),
      extra = Nothing
    }

resetIndex :: BH IO ()
resetIndex = do
  _ <- tryEsError deleteExampleIndex
  _ <- createExampleIndex
  _ <- performBHRequest $ putMapping @Value testIndex TweetMapping
  return ()

insertData :: BH IO (BHResponse StatusDependant IndexedDocument, IndexedDocument)
insertData = do
  _ <- tryEsError resetIndex
  insertData' defaultIndexDocumentSettings

insertData' :: IndexDocumentSettings -> BH IO (BHResponse StatusDependant IndexedDocument, IndexedDocument)
insertData' ids = do
  r <- performBHRequest $ keepBHResponse $ indexDocument testIndex ids exampleTweet (DocId "1")
  _ <- performBHRequest $ refreshIndex testIndex
  return r

-- | Returns the `BHResponse` of `indexDocument` without any parsing
insertData'' :: IndexDocumentSettings -> BH IO (BHResponse StatusDependant IndexedDocument)
insertData'' ids = do
  r <- dispatch $ indexDocument testIndex ids exampleTweet (DocId "1")
  _ <- performBHRequest $ refreshIndex testIndex
  return r

insertTweetWithDocId :: Tweet -> Text -> BH IO IndexedDocument
insertTweetWithDocId tweet docId = do
  let ids = defaultIndexDocumentSettings
  r <- performBHRequest $ indexDocument testIndex ids tweet (DocId docId)
  _ <- performBHRequest $ refreshIndex testIndex
  return r

updateData :: BH IO IndexedDocument
updateData = do
  r <- performBHRequest $ updateDocument testIndex defaultIndexDocumentSettings tweetPatch (DocId "1")
  _ <- performBHRequest $ refreshIndex testIndex
  return r

insertOther :: BH IO ()
insertOther = do
  _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings otherTweet (DocId "2")
  _ <- performBHRequest $ refreshIndex testIndex
  return ()

insertExtra :: BH IO ()
insertExtra = do
  _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings tweetWithExtra (DocId "4")
  _ <- performBHRequest $ refreshIndex testIndex
  return ()

insertWithSpaceInId :: BH IO ()
insertWithSpaceInId = do
  _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings exampleTweet (DocId "Hello World")
  _ <- performBHRequest $ refreshIndex testIndex
  return ()

searchTweet :: Search -> BH IO (Either EsError Tweet)
searchTweet search = (>>= grabFirst) <$> searchTweets search

searchTweets :: Search -> BH IO (Either EsError (SearchResult Tweet))
searchTweets search = tryPerformBHRequest $ searchByIndex testIndex search

searchExpectNoResults :: Search -> BH IO ()
searchExpectNoResults search = do
  result <- searchTweets search
  let emptyHits = fmap (hits . searchHits) result
  liftIO $
    emptyHits `shouldBe` Right []

searchExpectAggs :: Search -> BH IO ()
searchExpectAggs search = do
  result <- performBHRequest $ searchByIndex @Tweet testIndex search
  let isEmpty x = return (M.null x)
  liftIO $
    (aggregations result >>= isEmpty) `shouldBe` Just False

searchValidBucketAgg ::
  (BucketAggregation a, FromJSON a, Show a) =>
  Search ->
  Key ->
  (Key -> AggregationResults -> Maybe (Bucket a)) ->
  BH IO ()
searchValidBucketAgg search aggKey extractor = do
  result <- performBHRequest $ searchByIndex @Tweet testIndex search
  let bucketDocs = docCount . head . buckets
  let count = aggregations result >>= extractor aggKey >>= \x -> return (bucketDocs x)
  liftIO $
    count `shouldBe` Just 1

searchTermsAggHint :: [ExecutionHint] -> BH IO ()
searchTermsAggHint hints = do
  let terms hint = TermsAgg $ (mkTermsAggregation $ FieldName "user.keyword") {termExecutionHint = Just hint}
  let search hint = mkAggregateSearch Nothing $ mkAggregations "users" $ terms hint
  forM_ hints $ searchExpectAggs . search

-- forM_ hints (\x -> searchValidBucketAgg (search x) "users" toTerms)

searchTweetHighlight ::
  Search ->
  BH IO (Either EsError (Maybe HitHighlight))
searchTweetHighlight search = do
  result <- searchTweets search
  let tweetHit :: Either EsError (Maybe (Hit Tweet))
      tweetHit = fmap (headMay . hits . searchHits) result
      myHighlight :: Either EsError (Maybe HitHighlight)
      myHighlight = (join . fmap hitHighlight) <$> tweetHit
  return myHighlight

searchExpectSource :: Source -> Either EsError Value -> BH IO ()
searchExpectSource src expected = do
  _ <- insertData
  let query = QueryMatchQuery $ mkMatchQuery (FieldName "message") (QueryString "haskell")
  let search = (mkSearch (Just query) Nothing) {source = Just src}
  result <- performBHRequest $ searchByIndex testIndex search
  let value_ = grabFirst result
  liftIO $
    value_ `shouldBe` expected

is :: Versions.Version -> IO Bool
is v = getServerVersion >>= \x -> return $ x == v

esOnlyIT :: (HasCallStack, Example a) => IO (String -> a -> SpecWith (Arg a))
esOnlyIT = withMajorVersionIT (>= 6)

os2OnlyIT :: (HasCallStack, Example a) => IO (String -> a -> SpecWith (Arg a))
os2OnlyIT = withMajorVersionIT (== 2)

withMajorVersionIT :: (HasCallStack, Example a) => (Word -> Bool) -> IO (String -> a -> SpecWith (Arg a))
withMajorVersionIT p = do
  majoreVersion <- fetchMajorVersion
  return $
    if p majoreVersion
      then it
      else xit

fetchMajorVersion :: IO Word
fetchMajorVersion =
  withTestEnv $ do
    x <- performBHRequest $ getNodesInfo LocalNode
    let majoreVersion = versionNumber $ nodeInfoESVersion $ head $ nodesInfo x
    return $ head $ toListOf Versions.major majoreVersion