TestKit.hs

-- | A ready-made conformance test for a service that mounts this package's
-- probe routes.
--
-- Mounting @Servant.Health.HealthApi@ leaves a service one thing to prove: that
-- /its own/ assembled application really answers @GET \/health\/live@ and
-- @GET \/health\/ready@, that a failing readiness check produces a 503 while
-- liveness stays 200, and that the two routes were not wired to each other's
-- checks. That last hazard is the reason this kit exists. Both probe routes
-- have the same handler type, so an umbrella route record that swaps them
-- compiles cleanly and type-checks cleanly; only a request that distinguishes
-- them catches it.
--
-- 'probeContractTests' drives a real HTTP server built from the consumer's own
-- application and asserts the whole matrix, including the two cross-assertions
-- that fail when the routes are transposed.
--
-- == Wiring it into a service
--
-- Depend on the kit from the test stanza only, so production builds never
-- acquire its @tasty@ and @warp@ dependencies. The test suite must be built
-- with the threaded runtime: Warp's timer manager requires it, and without
-- @-threaded@ every request dies with @getSystemTimerManager: the TimerManager
-- requires linking against the threaded runtime@.
--
-- > test-suite myservice-test
-- >   type:           exitcode-stdio-1.0
-- >   hs-source-dirs: test
-- >   main-is:        Main.hs
-- >   ghc-options:    -threaded -rtsopts -with-rtsopts=-N
-- >   build-depends:
-- >     , base
-- >     , servant-health
-- >     , servant-health:testkit
-- >     , tasty
--
-- The kit supplies the two checks it wants mounted, so it needs a /builder/ for
-- the application rather than a finished one. Parameterizing the service's app
-- builder over its two probe checks is the only change most services need:
--
-- > module Main (main) where
-- >
-- > import GHC.Generics (Generic)
-- > import Network.Wai (Application)
-- > import Servant.API (NamedRoutes, (:>))
-- > import Servant.API.Generic ((:-))
-- > import Servant.Health (HealthApi, ProbeCheck, healthServer)
-- > import Servant.Health.TestKit (probeContractTests)
-- > import Servant.Server.Generic (genericServe)
-- > import Test.Tasty (defaultMain)
-- >
-- > -- Your service's umbrella route record. Normally this and 'mkApp' are
-- > -- imported from the service's own library rather than defined in the test
-- > -- file; they are inlined here so the example is complete.
-- > data MyServiceApi mode = MyServiceApi
-- >   { health :: mode :- "health" :> NamedRoutes HealthApi
-- >     -- , orders :: mode :- "orders" :> NamedRoutes OrdersApi
-- >   }
-- >   deriving stock (Generic)
-- >
-- > mkApp :: ProbeCheck -> ProbeCheck -> IO Application
-- > mkApp livenessCheck readinessCheck =
-- >   pure $
-- >     genericServe
-- >       MyServiceApi
-- >         { health = healthServer livenessCheck readinessCheck
-- >         }
-- >
-- > main :: IO ()
-- > main = defaultMain (probeContractTests "myservice probes" mkApp)
--
-- In production the same builder is called with the service's real checks,
-- assembled from "Servant.Health.Check". Note the nesting order: the timeout
-- goes /outside/ the exception wrapper, because
-- 'Servant.Health.Check.withProbeTimeout' interrupts by throwing an
-- asynchronous exception and 'Servant.Health.Check.safeCheck' deliberately lets
-- those through rather than swallowing them. Nested the other way round, a
-- deadlocked process would hang its liveness probe instead of failing it.
--
-- @ConnectionPool@, @eventLoopCheck@, @databaseCheck@, and @cacheCheck@ below
-- are the service's own — this package deliberately ships no checks:
--
-- > import Servant.Health.Check (safeCheck, sequenceChecks, withProbeTimeout)
-- >
-- > runService :: ConnectionPool -> IO ()
-- > runService pool = do
-- >   application <-
-- >     mkApp
-- >       (withProbeTimeout 2_000_000 "liveness" (safeCheck "liveness" (eventLoopCheck pool)))
-- >       (safeCheck "readiness" (sequenceChecks [databaseCheck pool, cacheCheck]))
-- >   Warp.run 8080 application
--
-- == What the builder must do
--
-- The builder must mount the checks it is handed, in the order it is handed
-- them — liveness first, readiness second — and it must mount them under
-- @"health"@ so the routes answer at 'Servant.Health.Paths.liveRawPath' and
-- 'Servant.Health.Paths.readyRawPath'. A builder that ignores its arguments and
-- mounts the service's real checks proves nothing and will fail these tests.
--
-- The builder is called once per test case, each time against a fresh server on
-- an ephemeral port, so cases are independent and safe to run in parallel. Keep
-- it cheap: a builder that opens a real database connection will open one per
-- case.
module Servant.Health.TestKit
  ( probeContractTests,
  )
where

import Control.Concurrent (rtsSupportsBoundThreads)
import Data.Aeson (Value, decode, object, (.=))
import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as Char8
import Data.ByteString.Lazy qualified as LazyByteString
import Data.IORef (newIORef, readIORef)
import Data.Text (Text)
import Data.Time (UTCTime (..), fromGregorian, secondsToDiffTime)
import Network.HTTP.Client
  ( Manager,
    Response,
    checkResponse,
    defaultManagerSettings,
    httpLbs,
    newManager,
    parseRequest,
    responseBody,
    responseHeaders,
    responseStatus,
  )
import Network.HTTP.Types.Header (hContentType)
import Network.HTTP.Types.Status (statusCode)
import Network.Wai (Application)
import Network.Wai.Handler.Warp (Port, testWithApplication)
import Servant.Health (ProbeCheck, ProbeVerdict (Healthy, Unhealthy))
import Servant.Health.Paths (liveRawPath, readyRawPath)
import Test.Tasty (TestName, TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertEqual, testCase)

-- | Prove that an assembled application answers the probe contract.
--
-- The second argument builds the application under test from a liveness check
-- and a readiness check, in that order. The kit owns those checks so that it
-- can decide their verdicts case by case; see the module documentation for a
-- complete example.
--
-- The returned group contains one case per scenario. Both checks healthy: both
-- probes answer 200 with the healthy body. Readiness failing: @\/health\/ready@
-- answers 503 naming the failed check while @\/health\/live@ still answers 200.
-- Liveness failing: the mirror image. The two \"still answers 200\" assertions
-- are the dispatch proof — they are what fails when an umbrella record wires
-- the probe routes to each other's checks.
probeContractTests ::
  TestName ->
  -- | Build the consumer's application from a liveness check and a readiness
  -- check. The kit owns the checks so it can flip them between cases.
  (ProbeCheck -> ProbeCheck -> IO Application) ->
  TestTree
probeContractTests groupName mkApp =
  testGroup
    groupName
    [ testCase "the test suite is built with the threaded runtime" $
        assertBool threadedRuntimeAdvice rtsSupportsBoundThreads,
      testCase "both checks healthy: both probes answer 200" $
        withProbeApp mkApp Healthy Healthy $ \manager port -> do
          getProbe manager port liveRawPath
            >>= assertProbe liveRawPath 200 healthyBody
          getProbe manager port readyRawPath
            >>= assertProbe readyRawPath 200 healthyBody,
      testCase "readiness failing: ready answers 503, live still 200" $
        withProbeApp mkApp Healthy (failing readinessCheckName) $ \manager port -> do
          getProbe manager port readyRawPath
            >>= assertProbe readyRawPath 503 (failingBody readinessCheckName)
          getProbe manager port liveRawPath
            >>= assertProbe liveRawPath 200 healthyBody,
      testCase "liveness failing: live answers 503, ready still 200" $
        withProbeApp mkApp (failing livenessCheckName) Healthy $ \manager port -> do
          getProbe manager port liveRawPath
            >>= assertProbe liveRawPath 503 (failingBody livenessCheckName)
          getProbe manager port readyRawPath
            >>= assertProbe readyRawPath 200 healthyBody
    ]
  where
    failing name = Unhealthy name probeFailureOnset

-- | What to say when the suite was linked without @-threaded@. Warp's timer
-- manager needs the threaded runtime, so without the flag every request in
-- every case below dies with a @getSystemTimerManager@ error that says nothing
-- about the missing flag.
threadedRuntimeAdvice :: String
threadedRuntimeAdvice =
  "This test suite was linked without the threaded runtime, so the probe "
    <> "server cannot run. Add -threaded -rtsopts -with-rtsopts=-N to the test "
    <> "suite's ghc-options in your .cabal file."

-- | The name the kit fails the liveness check under. Distinct from the
-- readiness name so a transposed wiring is named in the failure message rather
-- than merely counted.
livenessCheckName :: Text
livenessCheckName = "testkit-live"

-- | The name the kit fails the readiness check under.
readinessCheckName :: Text
readinessCheckName = "testkit-ready"

-- | The instant every failure the kit fabricates is dated with. Fixed rather
-- than read from the clock, so a failed assertion prints a stable expected
-- value.
probeFailureOnset :: UTCTime
probeFailureOnset =
  UTCTime (fromGregorian 2026 7 24) (secondsToDiffTime (12 * 3600 + 34 * 60 + 56))

-- | The body both probes answer with while their checks pass.
healthyBody :: Value
healthyBody =
  object
    [ "status" .= ("ok" :: Text),
      "check" .= ("all" :: Text),
      "failingSince" .= (Nothing :: Maybe UTCTime)
    ]

-- | The body a probe answers with when its check failed under the given name.
failingBody :: Text -> Value
failingBody failedCheckName =
  object
    [ "status" .= ("failed" :: Text),
      "check" .= failedCheckName,
      "failingSince" .= probeFailureOnset
    ]

-- | Build the consumer's application with checks that report the given
-- verdicts, serve it on an ephemeral port for the duration of the action, and
-- hand the action an HTTP client manager and the port.
--
-- The checks are backed by mutable cells read on every call, so an application
-- that caches a verdict at build time fails the cases above rather than passing
-- them by accident.
withProbeApp ::
  (ProbeCheck -> ProbeCheck -> IO Application) ->
  ProbeVerdict ->
  ProbeVerdict ->
  (Manager -> Port -> IO a) ->
  IO a
withProbeApp mkApp liveVerdict readyVerdict action = do
  liveRef <- newIORef liveVerdict
  readyRef <- newIORef readyVerdict
  application <- mkApp (readIORef liveRef) (readIORef readyRef)
  manager <- newManager defaultManagerSettings
  testWithApplication (pure application) (action manager)

-- | GET a raw probe path off the running server. The client's response check is
-- disabled so that a 503 arrives as an ordinary response rather than an
-- exception.
getProbe :: Manager -> Port -> ByteString -> IO (Response LazyByteString.ByteString)
getProbe manager port rawPath = do
  request <- parseRequest ("http://127.0.0.1:" <> show port <> Char8.unpack rawPath)
  httpLbs request {checkResponse = \_ _ -> pure ()} manager

-- | Assert the status code, JSON body, and content type of one probe response,
-- labelling each assertion with the path it was made against so a failure names
-- the probe rather than just the expected value.
assertProbe ::
  ByteString ->
  Int ->
  Value ->
  Response LazyByteString.ByteString ->
  IO ()
assertProbe rawPath expectedStatus expectedBody response = do
  assertEqual
    (label "status code")
    expectedStatus
    (statusCode (responseStatus response))
  assertEqual
    (label "JSON body")
    (Just expectedBody)
    (decode (responseBody response))
  assertEqual
    (label "content type")
    (Just "application/json")
    (mediaType (lookup hContentType (responseHeaders response)))
  where
    label what = "GET " <> Char8.unpack rawPath <> ": " <> what
    mediaType = fmap (Char8.takeWhile (/= ';'))