{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

module Test.Sandwich.Contexts.Kubernetes.Waits where

import Control.Monad.IO.Unlift
import Control.Monad.Logger
import qualified Data.List as L
import qualified Data.Map as M
import Data.String.Interpolate
import qualified Data.Text as T
import Kubernetes.OpenAPI.API.CoreV1 as Kubernetes
import Kubernetes.OpenAPI.Core as Kubernetes
import Kubernetes.OpenAPI.MimeTypes
import Kubernetes.OpenAPI.Model as Kubernetes
import Relude
import System.Exit
import Test.Sandwich
import Test.Sandwich.Contexts.Files
import Test.Sandwich.Contexts.Kubernetes.Run
import Test.Sandwich.Contexts.Kubernetes.Types
import Test.Sandwich.Waits
import UnliftIO.Process


-- | Wait for a service to have its set of endpoints ready, i.e.:
--
-- * They each have at least one IP address
-- * They each have an empty set of "not ready addresses"
waitForServiceEndpointsToExist :: (
  MonadUnliftIO m, MonadLogger m
  , MonadReader context m, HasKubernetesClusterContext context
  )
  -- | Namespace
  => Text
  -- | Service name
  -> Text
  -- | Time in seconds to wait
  -> Double
  -> m ()
waitForServiceEndpointsToExist :: forall (m :: * -> *) context.
(MonadUnliftIO m, MonadLogger m, MonadReader context m,
 HasKubernetesClusterContext context) =>
Text -> Text -> Double -> m ()
waitForServiceEndpointsToExist Text
namespace Text
serviceName Double
timeInSeconds = do
  Double -> m () -> m ()
forall (m :: * -> *) a.
(HasCallStack, MonadUnliftIO m) =>
Double -> m a -> m a
waitUntil Double
timeInSeconds (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    endpoints <- Text -> Map Text Text -> m [V1Endpoints]
forall (m :: * -> *) context.
(MonadUnliftIO m, MonadLogger m, MonadReader context m,
 HasKubernetesClusterContext context) =>
Text -> Map Text Text -> m [V1Endpoints]
listEndpoints Text
namespace Map Text Text
forall a. Monoid a => a
mempty
    case Relude.filter v1EndpointsSatisfies endpoints of
      [] -> String -> m ()
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|No endpoints were satisfactory|]
      (V1Endpoints
x:[V1Endpoints]
_) -> do
        Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
debug [i|(#{namespace}) Got satisfactory endpoint for #{serviceName}: #{x}|]
        () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

  where
    v1EndpointsSatisfies :: V1Endpoints -> Bool
v1EndpointsSatisfies (V1Endpoints {v1EndpointsMetadata :: V1Endpoints -> Maybe V1ObjectMeta
v1EndpointsMetadata=(Just (V1ObjectMeta {v1ObjectMetaName :: V1ObjectMeta -> Maybe Text
v1ObjectMetaName=(Just Text
name)})), Maybe [V1EndpointSubset]
v1EndpointsSubsets :: Maybe [V1EndpointSubset]
v1EndpointsSubsets :: V1Endpoints -> Maybe [V1EndpointSubset]
v1EndpointsSubsets})
      | Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
serviceName = (V1EndpointSubset -> Bool) -> [V1EndpointSubset] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
L.all V1EndpointSubset -> Bool
isSatisfactoryV1EndpointSubset ([V1EndpointSubset]
-> Maybe [V1EndpointSubset] -> [V1EndpointSubset]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [V1EndpointSubset]
v1EndpointsSubsets)
    v1EndpointsSatisfies V1Endpoints
_ = Bool
False

    isSatisfactoryV1EndpointSubset :: V1EndpointSubset -> Bool
isSatisfactoryV1EndpointSubset (V1EndpointSubset {
                                       v1EndpointSubsetAddresses :: V1EndpointSubset -> Maybe [V1EndpointAddress]
v1EndpointSubsetAddresses=(Just [V1EndpointAddress]
addrs)
                                       , v1EndpointSubsetNotReadyAddresses :: V1EndpointSubset -> Maybe [V1EndpointAddress]
v1EndpointSubsetNotReadyAddresses=([V1EndpointAddress]
-> Maybe [V1EndpointAddress] -> [V1EndpointAddress]
forall a. a -> Maybe a -> a
fromMaybe [] -> [V1EndpointAddress]
notReadyAddrs)
                                       }) =
      Bool -> Bool
not ([V1EndpointAddress] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
L.null [V1EndpointAddress]
addrs)
      Bool -> Bool -> Bool
&& [V1EndpointAddress] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
L.null [V1EndpointAddress]
notReadyAddrs
    isSatisfactoryV1EndpointSubset V1EndpointSubset
_ = Bool
False


listEndpoints :: (
  MonadUnliftIO m, MonadLogger m
  , MonadReader context m, HasKubernetesClusterContext context
  ) => Text -> Map Text Text -> m [V1Endpoints]
listEndpoints :: forall (m :: * -> *) context.
(MonadUnliftIO m, MonadLogger m, MonadReader context m,
 HasKubernetesClusterContext context) =>
Text -> Map Text Text -> m [V1Endpoints]
listEndpoints Text
namespace Map Text Text
labels =
  (V1EndpointsList -> [V1Endpoints]
v1EndpointsListItems (V1EndpointsList -> [V1Endpoints])
-> m V1EndpointsList -> m [V1Endpoints]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (m V1EndpointsList -> m [V1Endpoints])
-> m V1EndpointsList -> m [V1Endpoints]
forall a b. (a -> b) -> a -> b
$ KubernetesRequest
  ListNamespacedEndpoints MimeNoContent V1EndpointsList MimeJSON
-> m V1EndpointsList
forall req accept res contentType context (m :: * -> *).
(Produces req accept, MimeUnrender accept res,
 MimeType contentType, Constraints context m,
 HasKubernetesClusterContext context) =>
KubernetesRequest req contentType res accept -> m res
k8sRunException (
      (Accept MimeJSON
-> Namespace
-> KubernetesRequest
     ListNamespacedEndpoints MimeNoContent V1EndpointsList MimeJSON
forall accept.
Accept accept
-> Namespace
-> KubernetesRequest
     ListNamespacedEndpoints MimeNoContent V1EndpointsList accept
listNamespacedEndpoints (MimeJSON -> Accept MimeJSON
forall a. MimeType a => a -> Accept a
Accept MimeJSON
MimeJSON) (Text -> Namespace
Namespace Text
namespace))
      KubernetesRequest
  ListNamespacedEndpoints MimeNoContent V1EndpointsList MimeJSON
-> LabelSelector
-> KubernetesRequest
     ListNamespacedEndpoints MimeNoContent V1EndpointsList MimeJSON
forall req param contentType res accept.
HasOptionalParam req param =>
KubernetesRequest req contentType res accept
-> param -> KubernetesRequest req contentType res accept
forall contentType res accept.
KubernetesRequest ListNamespacedEndpoints contentType res accept
-> LabelSelector
-> KubernetesRequest ListNamespacedEndpoints contentType res accept
-&- (Text -> LabelSelector
LabelSelector (Text -> [Text] -> Text
T.intercalate Text
"," [Text
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"=" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
v | (Text
k, Text
v) <- Map Text Text -> [(Text, Text)]
forall k a. Map k a -> [(k, a)]
M.toList Map Text Text
labels]))
    )

-- | Wait for a set of pods to exist, specified by a set of labels.
waitForPodsToExist :: (
  MonadUnliftIO m, MonadLogger m
  , MonadReader context m, HasKubernetesClusterContext context
  )
  -- | Namespace
  => Text
  -- | Pod labels
  -> Map Text Text
  -- | Time in seconds to wait
  -> Double
  -- | Optional desired pod count to wait for
  -> Maybe Int
  -> m ()
waitForPodsToExist :: forall (m :: * -> *) context.
(MonadUnliftIO m, MonadLogger m, MonadReader context m,
 HasKubernetesClusterContext context) =>
Text -> Map Text Text -> Double -> Maybe Int -> m ()
waitForPodsToExist Text
namespace Map Text Text
labels Double
timeInSeconds Maybe Int
maybeDesiredCount = do
  Double -> m () -> m ()
forall (m :: * -> *) a.
(HasCallStack, MonadUnliftIO m) =>
Double -> m a -> m a
waitUntil Double
timeInSeconds (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    pods <- Text -> Map Text Text -> m [V1Pod]
forall (m :: * -> *) context.
(MonadUnliftIO m, MonadLogger m, MonadReader context m,
 HasKubernetesClusterContext context) =>
Text -> Map Text Text -> m [V1Pod]
listPods Text
namespace Map Text Text
labels
    case maybeDesiredCount of
      Maybe Int
Nothing -> Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([V1Pod] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
L.null [V1Pod]
pods) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Found no pods.|]
      Just Int
n -> Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([V1Pod] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
L.length [V1Pod]
pods Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Expected #{n} pods, but found #{L.length pods}|]

-- | List the pods matching a set of labels.
listPods :: (
  MonadUnliftIO m, MonadLogger m
  , MonadReader context m, HasKubernetesClusterContext context
  ) => Text -> Map Text Text -> m [V1Pod]
listPods :: forall (m :: * -> *) context.
(MonadUnliftIO m, MonadLogger m, MonadReader context m,
 HasKubernetesClusterContext context) =>
Text -> Map Text Text -> m [V1Pod]
listPods Text
namespace Map Text Text
labels =
  (V1PodList -> [V1Pod]
v1PodListItems (V1PodList -> [V1Pod]) -> m V1PodList -> m [V1Pod]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>) (m V1PodList -> m [V1Pod]) -> m V1PodList -> m [V1Pod]
forall a b. (a -> b) -> a -> b
$ KubernetesRequest
  ListNamespacedPod MimeNoContent V1PodList MimeJSON
-> m V1PodList
forall req accept res contentType context (m :: * -> *).
(Produces req accept, MimeUnrender accept res,
 MimeType contentType, Constraints context m,
 HasKubernetesClusterContext context) =>
KubernetesRequest req contentType res accept -> m res
k8sRunException (
      (Accept MimeJSON
-> Namespace
-> KubernetesRequest
     ListNamespacedPod MimeNoContent V1PodList MimeJSON
forall accept.
Accept accept
-> Namespace
-> KubernetesRequest
     ListNamespacedPod MimeNoContent V1PodList accept
listNamespacedPod (MimeJSON -> Accept MimeJSON
forall a. MimeType a => a -> Accept a
Accept MimeJSON
MimeJSON) (Text -> Namespace
Namespace Text
namespace))
      KubernetesRequest
  ListNamespacedPod MimeNoContent V1PodList MimeJSON
-> LabelSelector
-> KubernetesRequest
     ListNamespacedPod MimeNoContent V1PodList MimeJSON
forall req param contentType res accept.
HasOptionalParam req param =>
KubernetesRequest req contentType res accept
-> param -> KubernetesRequest req contentType res accept
forall contentType res accept.
KubernetesRequest ListNamespacedPod contentType res accept
-> LabelSelector
-> KubernetesRequest ListNamespacedPod contentType res accept
-&- (Text -> LabelSelector
LabelSelector (Text -> [Text] -> Text
T.intercalate Text
"," [Text
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"=" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
v | (Text
k, Text
v) <- Map Text Text -> [(Text, Text)]
forall k a. Map k a -> [(k, a)]
M.toList Map Text Text
labels]))
    )

-- | Wait for a set of pods to be in the Ready condition, specified by a set of labels.
waitForPodsToBeReady :: (
  MonadUnliftIO m, MonadLogger m
  , MonadReader context m, HasKubernetesClusterContext context, HasFile context "kubectl"
  )
  -- | Namespace
  => Text
  -- | Pod labels
  -> Map Text Text
  -- | Time in seconds to wait
  -> Double
  -> m ()
waitForPodsToBeReady :: forall (m :: * -> *) context.
(MonadUnliftIO m, MonadLogger m, MonadReader context m,
 HasKubernetesClusterContext context, HasFile context "kubectl") =>
Text -> Map Text Text -> Double -> m ()
waitForPodsToBeReady Text
namespace Map Text Text
labels Double
timeInSeconds = do
  kubectlBinary <- forall (a :: Symbol) context (m :: * -> *).
(MonadReader context m, HasFile context a) =>
m String
askFile @"kubectl"
  kubeConfigFile <- kubernetesClusterKubeConfigPath <$> getContext kubernetesCluster

  let labelArgs = [[i|-l #{k}=#{v}|] | (Text
k, Text
v) <- Map Text Text -> [(Text, Text)]
forall k a. Map k a -> [(k, a)]
M.toList Map Text Text
labels]
  p <- createProcessWithLogging (proc kubectlBinary (
                                  ["wait", "pods"
                                  , "--kubeconfig", kubeConfigFile
                                  , "-n", toString namespace
                                  ]
                                  <> labelArgs
                                  <> [
                                    "--for", "condition=Ready"
                                    , "--timeout=" <> show timeInSeconds <> "s"
                                    ]
                                ))
  waitForProcess p >>= \case
    ExitCode
ExitSuccess -> () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    ExitFailure Int
n -> String -> m ()
forall (m :: * -> *) a. (HasCallStack, MonadIO m) => String -> m a
expectationFailure [i|Failed to wait for pods to exist (code #{n})|]