AggregationSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.AggregationSpec (spec) where

import Control.Error (fmapL, note)
import qualified Data.Map as M
import qualified Database.Bloodhound
import TestsUtils.Common
import TestsUtils.Import

spec :: Spec
spec =
  describe "Aggregation API" $ do
    it "returns term aggregation results" $
      withTestEnv $ do
        _ <- insertData
        let terms = TermsAgg $ mkTermsAggregation "user"
        let search = mkAggregateSearch Nothing $ mkAggregations "users" terms
        searchExpectAggs search
        searchValidBucketAgg search "users" toTerms

    it "return sub-aggregation results" $
      withTestEnv $ do
        _ <- insertData
        let subaggs = mkAggregations "age_agg" . TermsAgg $ mkTermsAggregation "age"
            agg = TermsAgg $ (mkTermsAggregation "user") {termAggs = Just subaggs}
            search = mkAggregateSearch Nothing $ mkAggregations "users" agg
        result <- performBHRequest $ searchByIndex @Tweet testIndex search
        let usersAggResults = aggregations result >>= toTerms "users"
            subAggResults = usersAggResults >>= (listToMaybe . buckets) >>= termsAggs >>= toTerms "age_agg"
            subAddResultsExists = isJust subAggResults
        liftIO $ subAddResultsExists `shouldBe` True

    it "returns cardinality aggregation results" $
      withTestEnv $ do
        _ <- insertData
        let cardinality = CardinalityAgg $ mkCardinalityAggregation $ FieldName "user"
        let search = mkAggregateSearch Nothing $ mkAggregations "users" cardinality
        let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
        searchExpectAggs search'
        let docCountPair k n = (k, object ["value" .= Number n])
        res <- searchTweets search'
        liftIO $
          fmap aggregations res `shouldBe` Right (Just (M.fromList [docCountPair "users" 1]))

    it "returns stats aggregation results" $
      withTestEnv $ do
        _ <- insertData
        let stats = StatsAgg $ mkStatsAggregation $ FieldName "age"
        let search = mkAggregateSearch Nothing $ mkAggregations "users" stats
        let search' = search {Database.Bloodhound.from = From 0, size = Size 0}
        searchExpectAggs search'
        let statsAggRes k n =
              ( k,
                object
                  [ "max" .= Number n,
                    "avg" .= Number n,
                    "count" .= Number 1,
                    "min" .= Number n,
                    "sum" .= Number n
                  ]
              )
        res <- searchTweets search'
        liftIO $
          fmap aggregations res `shouldBe` Right (Just (M.fromList [statsAggRes "users" 10000]))

    it "can give collection hint parameters to term aggregations" $
      withTestEnv $ do
        _ <- insertData
        let terms = TermsAgg $ (mkTermsAggregation "user") {termCollectMode = Just BreadthFirst}
        let search = mkAggregateSearch Nothing $ mkAggregations "users" terms
        searchExpectAggs search
        searchValidBucketAgg search "users" toTerms

    it "can give execution hint parameters to term aggregations" $
      withTestEnv $ do
        _ <- insertData
        searchTermsAggHint [GlobalOrdinals, Map]
    -- One of the above.

    it "can execute value_count aggregations" $
      withTestEnv $ do
        _ <- insertData
        _ <- insertOther
        let ags =
              mkAggregations "user_count" (ValueCountAgg (FieldValueCount (FieldName "user")))
                <> mkAggregations "bogus_count" (ValueCountAgg (FieldValueCount (FieldName "bogus")))
        let search = mkAggregateSearch Nothing ags
        let docCountPair k n = (k, object ["value" .= Number n])
        res <- searchTweets search
        liftIO $
          fmap aggregations res
            `shouldBe` Right
              ( Just
                  ( M.fromList
                      [ docCountPair "user_count" 2,
                        docCountPair "bogus_count" 0
                      ]
                  )
              )

    it "can execute date_range aggregations" $
      withTestEnv $ do
        let now = fromGregorian 2015 3 14
        let ltAMonthAgo = UTCTime (fromGregorian 2015 3 1) 0
        let ltAWeekAgo = UTCTime (fromGregorian 2015 3 10) 0
        let oldDoc = exampleTweet {postDate = ltAMonthAgo}
        let newDoc = exampleTweet {postDate = ltAWeekAgo}
        _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings oldDoc (DocId "1")
        _ <- performBHRequest $ indexDocument testIndex defaultIndexDocumentSettings newDoc (DocId "2")
        _ <- performBHRequest $ refreshIndex testIndex
        let thisMonth = DateRangeFrom (DateMathExpr (DMDate now) [SubtractTime 1 DMMonth])
        let thisWeek = DateRangeFrom (DateMathExpr (DMDate now) [SubtractTime 1 DMWeek])
        let agg = DateRangeAggregation (FieldName "postDate") Nothing (thisMonth :| [thisWeek])
        let ags = mkAggregations "date_ranges" (DateRangeAgg agg)
        let search = mkAggregateSearch Nothing ags
        res <- searchTweets search
        liftIO $ hitsTotal . searchHits <$> res `shouldBe` Right (HitsTotal 2 HTR_EQ)
        let bucks = do
              magrs <- fmapL show (aggregations <$> res)
              agrs <- note "no aggregations returned" magrs
              rawBucks <- note "no date_ranges aggregation" $ M.lookup "date_ranges" agrs
              parseEither parseJSON rawBucks
        let fromMonthT = UTCTime (fromGregorian 2015 2 14) 0
        let fromWeekT = UTCTime (fromGregorian 2015 3 7) 0
        liftIO $
          buckets <$> bucks
            `shouldBe` Right
              [ DateRangeResult
                  "2015-02-14T00:00:00.000Z-*"
                  (Just fromMonthT)
                  (Just "2015-02-14T00:00:00.000Z")
                  Nothing
                  Nothing
                  2
                  Nothing,
                DateRangeResult
                  "2015-03-07T00:00:00.000Z-*"
                  (Just fromWeekT)
                  (Just "2015-03-07T00:00:00.000Z")
                  Nothing
                  Nothing
                  1
                  Nothing
              ]

    it "returns date histogram aggregation results" $
      withTestEnv $ do
        _ <- insertData
        let histogram = DateHistogramAgg $ mkDateHistogram (FieldName "postDate") Minute
        let search = mkAggregateSearch Nothing (mkAggregations "byDate" histogram)
        searchExpectAggs search
        searchValidBucketAgg search "byDate" toDateHistogram

    it "can execute missing aggregations" $
      withTestEnv $ do
        _ <- insertData
        _ <- insertExtra
        let ags = mkAggregations "missing_agg" (MissingAgg (MissingAggregation "extra"))
        let search = mkAggregateSearch Nothing ags
        let docCountPair k n = (k, object ["doc_count" .= Number n])
        res <- searchTweets search
        liftIO $
          fmap aggregations res `shouldBe` Right (Just (M.fromList [docCountPair "missing_agg" 1]))