{-# LANGUAGE OverloadedStrings #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Test.Federator where

import API.Brig
import API.Federator (getMetrics)
import Data.Attoparsec.Text
import qualified Data.ByteString as BS
import Data.String.Conversions
import Data.Text
import SetupHelpers (randomUser)
import Testlib.Prelude

runFederatorMetrics :: (ServiceMap -> HostPort) -> App ()
runFederatorMetrics :: (ServiceMap -> HostPort) -> App ()
runFederatorMetrics ServiceMap -> HostPort
getServiceAddress = do
  let handleRes :: r -> App r
handleRes r
res = r
res r -> App () -> App r
forall a b. a -> App b -> App a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ r
res.status a -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
  first <- App Response -> (Response -> App Response) -> App Response
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
bindResponse (Domain -> (ServiceMap -> HostPort) -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> (ServiceMap -> HostPort) -> App Response
getMetrics Domain
OwnDomain ServiceMap -> HostPort
getServiceAddress) Response -> App Response
forall {a} {r}. (MakesValue a, HasField "status" r a) => r -> App r
handleRes
  second <- bindResponse (getMetrics OwnDomain getServiceAddress) handleRes
  assertBool "Two metric requests should never match" $ first.body /= second.body
  assertBool "Second metric response should never be 0 length (the first might be)" $ BS.length second.body /= 0
  assertBool "The seconds metric response should have text indicating that it is returning metrics"
    $ BS.isInfixOf expectedString second.body
  where
    expectedString :: ByteString
expectedString = ByteString
"# TYPE http_request_duration_seconds histogram"

-- | The metrics setup for both internal and external federator servers
-- are the same, so we can simply run the same test for both.
testFederatorMetricsInternal :: App ()
testFederatorMetricsInternal :: App ()
testFederatorMetricsInternal = (ServiceMap -> HostPort) -> App ()
runFederatorMetrics ServiceMap -> HostPort
federatorInternal

testFederatorMetricsExternal :: App ()
testFederatorMetricsExternal :: App ()
testFederatorMetricsExternal = (ServiceMap -> HostPort) -> App ()
runFederatorMetrics ServiceMap -> HostPort
federatorExternal

testFederatorNumRequestsMetrics :: (HasCallStack) => App ()
testFederatorNumRequestsMetrics :: HasCallStack => App ()
testFederatorNumRequestsMetrics = do
  u1 <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  u2 <- randomUser OtherDomain def
  incomingBefore <- getMetric parseIncomingRequestCount OtherDomain OwnDomain
  outgoingBefore <- getMetric parseOutgoingRequestCount OwnDomain OtherDomain
  bindResponse (searchContacts u1 (u2 %. "name") OtherDomain) $ \Response
resp ->
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
  incomingAfter <- getMetric parseIncomingRequestCount OtherDomain OwnDomain
  outgoingAfter <- getMetric parseOutgoingRequestCount OwnDomain OtherDomain
  assertBool "Incoming requests count should have increased by at least 2" $ incomingAfter >= incomingBefore + 2
  assertBool "Outgoing requests count should have increased by at least 2" $ outgoingAfter >= outgoingBefore + 2
  where
    getMetric :: (Text -> Parser Integer) -> Domain -> Domain -> App Integer
    getMetric :: (Text -> Parser Integer) -> Domain -> Domain -> App Integer
getMetric Text -> Parser Integer
p Domain
domain Domain
origin = do
      m <- Domain -> (ServiceMap -> HostPort) -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> (ServiceMap -> HostPort) -> App Response
getMetrics Domain
domain ServiceMap -> HostPort
federatorInternal
      d <- cs <$> asString origin
      pure $ fromRight 0 (parseOnly (p d) (cs m.body))

    parseIncomingRequestCount :: Text -> Parser Integer
    parseIncomingRequestCount :: Text -> Parser Integer
parseIncomingRequestCount Text
d =
      Parser Text Char -> Parser Text Text -> Parser Text String
forall (f :: * -> *) a b. Alternative f => f a -> f b -> f [a]
manyTill Parser Text Char
anyChar (Text -> Parser Text Text
string (Text
"com_wire_federator_incoming_requests{origin_domain=\"" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
d Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\"} "))
        Parser Text String -> Parser Integer -> Parser Integer
forall a b. Parser Text a -> Parser Text b -> Parser Text b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Parser Integer
forall a. Integral a => Parser a
decimal

    parseOutgoingRequestCount :: Text -> Parser Integer
    parseOutgoingRequestCount :: Text -> Parser Integer
parseOutgoingRequestCount Text
d =
      Parser Text Char -> Parser Text Text -> Parser Text String
forall (f :: * -> *) a b. Alternative f => f a -> f b -> f [a]
manyTill Parser Text Char
anyChar (Text -> Parser Text Text
string (Text
"com_wire_federator_outgoing_requests{target_domain=\"" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
d Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\"} "))
        Parser Text String -> Parser Integer -> Parser Integer
forall a b. Parser Text a -> Parser Text b -> Parser Text b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Parser Integer
forall a. Integral a => Parser a
decimal