DiskUsageSpec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Test.DiskUsageSpec (spec) where

import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy qualified as LBS
import Data.List (sort)
import Database.Bloodhound.Common.Types
import Optics (set, view)
import TestsUtils.Common
import TestsUtils.Import

-- | Stable ordering for equality.
normalize :: [(Text, Maybe Text)] -> [(Text, Maybe Text)]
normalize = sort

-- | A trimmed but faithful slice of the documented response: @\"_shards\"@
-- plus a single index key carrying store sizes, an @all_fields@ aggregate
-- and one per-field entry.
sampleResponseBytes :: LBS.ByteString
sampleResponseBytes =
  "{\"_shards\":{\"total\":1,\"successful\":1,\"failed\":0},\
  \\"my-index-000001\":{\
  \\"store_size\":\"929mb\",\
  \\"store_size_in_bytes\":974192723,\
  \\"all_fields\":{\"total\":\"928.9mb\",\"total_in_bytes\":973977084,\
  \\"inverted_index\":{\"total\":\"107.8mb\",\"total_in_bytes\":113128526}},\
  \\"fields\":{\"_id\":{\"total\":\"49.3mb\",\"total_in_bytes\":51709993,\
  \\"doc_values\":\"0b\",\"doc_values_in_bytes\":0}}}}"

spec :: Spec
spec = do
  -- ------------------------------------------------------------------ --
  -- Pure JSON decoder tests (no ES required).                          --
  -- ------------------------------------------------------------------ --
  describe "DiskUsageResponse JSON parsing" $ do
    it "peels _shards and keeps the rest as a per-index KeyMap" $ do
      let Just (r :: DiskUsageResponse) = decode sampleResponseBytes
      shardTotal (durShards r) `shouldBe` 1
      shardsSuccessful (durShards r) `shouldBe` 1
      case KM.lookup "my-index-000001" (durIndices r) of
        Just ix -> do
          duiStoreSize ix `shouldBe` Just "929mb"
          duiStoreSizeInBytes ix `shouldBe` Just 974192723
          case duiAllFields ix of
            Just af -> do
              dufbTotal af `shouldBe` Just "928.9mb"
              dufbTotalInBytes af `shouldBe` Just 973977084
              case dufbInvertedIndex af of
                Just ii -> do
                  duiiTotal ii `shouldBe` Just "107.8mb"
                  duiiTotalInBytes ii `shouldBe` Just 113128526
                Nothing -> expectationFailure "expected inverted_index"
            Nothing -> expectationFailure "expected all_fields"
          case duiFields ix of
            Just fields -> case KM.lookup "_id" fields of
              Just fb -> do
                dufbTotalInBytes fb `shouldBe` Just 51709993
                dufbDocValuesInBytes fb `shouldBe` Just 0
              Nothing -> expectationFailure "expected _id field"
            Nothing -> expectationFailure "expected fields map"
        Nothing -> expectationFailure "expected my-index-000001 entry"

    it "round-trips the documented response via ToJSON/FromJSON" $ do
      let Just (original :: DiskUsageResponse) = decode sampleResponseBytes
          reDecoded = decode (encode original) :: Maybe DiskUsageResponse
      Just original `shouldBe` reDecoded

    it "tolerates a response with only _shards (empty index map)" $ do
      let Just (r :: DiskUsageResponse) =
            decode "{\"_shards\":{\"total\":0,\"successful\":0,\"failed\":0}}"
      durIndices r `shouldBe` mempty

  describe "DiskUsageFieldBreakdown JSON parsing" $ do
    let bare = "{\"total\":\"1b\",\"total_in_bytes\":1}"
    it "parses a minimal breakdown (no structural components)" $ do
      let Just (fb :: DiskUsageFieldBreakdown) = decode bare
      dufbTotal fb `shouldBe` Just "1b"
      dufbTotalInBytes fb `shouldBe` Just 1
      dufbInvertedIndex fb `shouldBe` Nothing
      dufbStoredFieldsInBytes fb `shouldBe` Nothing

    it "round-trips via ToJSON/FromJSON" $ do
      let Just (original :: DiskUsageFieldBreakdown) = decode bare
          reDecoded = decode (encode original) :: Maybe DiskUsageFieldBreakdown
      Just original `shouldBe` reDecoded

  -- ------------------------------------------------------------------ --
  -- Optics.                                                            --
  -- ------------------------------------------------------------------ --
  describe "DiskUsage lenses" $ do
    it "duiStoreSizeInBytesLens is a lawful get/set" $ do
      let ix = DiskUsageIndex {duiStoreSize = Nothing, duiStoreSizeInBytes = Nothing, duiAllFields = Nothing, duiFields = Nothing}
          ix' = set duiStoreSizeInBytesLens (Just 42) ix
      view duiStoreSizeInBytesLens ix' `shouldBe` Just 42

    it "duoRunExpensiveTasksLens is a lawful get/set" $ do
      let opts = set duoRunExpensiveTasksLens False defaultDiskUsageOptions
      view duoRunExpensiveTasksLens opts `shouldBe` False

  -- ------------------------------------------------------------------ --
  -- Options rendering.                                                 --
  -- ------------------------------------------------------------------ --
  describe "DiskUsageOptions URI param rendering" $ do
    it "defaultDiskUsageOptions emits only run_expensive_tasks=true" $
      diskUsageOptionsParams defaultDiskUsageOptions
        `shouldBe` [("run_expensive_tasks", Just "true")]

    it "always renders run_expensive_tasks, even when False" $ do
      let opts = defaultDiskUsageOptions {duoRunExpensiveTasks = False}
      lookup "run_expensive_tasks" (diskUsageOptionsParams opts)
        `shouldBe` Just (Just "false")

    it "renders the optional params when set" $ do
      let opts =
            defaultDiskUsageOptions
              { duoFlush = Just False,
                duoIgnoreUnavailable = Just True,
                duoExpandWildcards = Just (ExpandWildcardsOpen :| [])
              }
      normalize (diskUsageOptionsParams opts)
        `shouldBe` [ ("expand_wildcards", Just "open"),
                     ("flush", Just "false"),
                     ("ignore_unavailable", Just "true"),
                     ("run_expensive_tasks", Just "true")
                   ]

  -- ------------------------------------------------------------------ --
  -- Endpoint shape (BHRequest, no ES required).                        --
  -- ------------------------------------------------------------------ --
  describe "diskUsage endpoint shape" $ do
    let idx = [qqIndexName|my-idx|]

    it "POSTs /{index}/_disk_usage with run_expensive_tasks=true by default" $ do
      let req = diskUsage idx
      bhRequestMethod req `shouldBe` "POST"
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["my-idx", "_disk_usage"]
      getRawEndpointQueries (bhRequestEndpoint req)
        `shouldBe` [("run_expensive_tasks", Just "true")]

    it "sends an empty body" $ do
      let req = diskUsage idx
      bhRequestBody req `shouldBe` Just ""

    it "forwards run_expensive_tasks=false via the With variant" $ do
      let opts = defaultDiskUsageOptions {duoRunExpensiveTasks = False}
          req = diskUsageWith opts idx
      lookup "run_expensive_tasks" (getRawEndpointQueries (bhRequestEndpoint req))
        `shouldBe` Just (Just "false")