Generator.hs

{-# LANGUAGE OverloadedLists #-}

-- |
-- Module:      Data.OpenApi.Schema.Generator
-- Maintainer:  Nadeem Bitar <nadeem@gmail.com>
-- Stability:   experimental
--
-- QuickCheck generators that produce JSON 'Value's conforming to an OpenAPI
-- 3.1 'Schema'. This is primarily useful for property testing: derive a schema
-- from a type with 'ToSchema', generate random values that satisfy it, and
-- check that they parse back via 'FromJSON' (see 'validateFromJSON').
module Data.OpenApi.Schema.Generator where

import Control.Lens.Operators
import Control.Monad (filterM)
import Data.Aeson
import Data.Aeson.Types
import Data.HashMap.Strict.InsOrd.Compat qualified as M
import Data.Maybe
import Data.OpenApi
import Data.OpenApi.Aeson.Compat (fromInsOrdHashMap)
import Data.OpenApi.Declare
import Data.OpenApi.Internal.Schema.Validation (inferSchemaTypes)
import Data.Proxy
import Data.Scientific
import Data.Set qualified as S
import Data.Text qualified as T
import Data.Vector qualified as V
import Prelude.Compat
import Test.QuickCheck (arbitrary)
import Test.QuickCheck.Gen
import Test.QuickCheck.Property
import Prelude ()

-- | Note: 'schemaGen' may 'error', if schema type is not specified,
-- and cannot be inferred.
schemaGen :: Definitions Schema -> Schema -> Gen Value
schemaGen _ schema
  | Just cases <- schema ^. enum_ = elements cases
schemaGen defns schema
  | Just variants <- schema ^. oneOf = schemaGen defns =<< elements (dereference defns <$> variants)
schemaGen defns schema =
  case singleType =<< schema ^. type_ of
    Nothing ->
      case inferSchemaTypes schema of
        [inferredType] -> schemaGen defns (schema & type_ ?~ OpenApiTypeSingle inferredType)
        -- Gen is not MonadFail
        _ -> error "unable to infer schema type"
    Just OpenApiBoolean -> Bool <$> elements [True, False]
    Just OpenApiNull -> pure Null
    Just OpenApiNumber
      | Just min <- schema ^. minimum_,
        Just max <- schema ^. maximum_ ->
          Number
            . fromFloatDigits
            <$> choose (toRealFloat min, toRealFloat max :: Double)
      | otherwise -> Number . fromFloatDigits <$> (arbitrary :: Gen Double)
    Just OpenApiInteger
      | Just min <- schema ^. minimum_,
        Just max <- schema ^. maximum_ ->
          Number
            . fromInteger
            <$> choose (truncate min, truncate max)
      | otherwise -> Number . fromInteger <$> arbitrary
    Just OpenApiArray
      | Just 0 <- schema ^. maxLength -> pure $ Array V.empty
      | Just prefixSchemas <- schema ^. prefixItems -> do
          -- 3.1 tuple: generate one element positionally per prefix schema, then
          -- any trailing elements permitted by `items` (for the generic tuple
          -- derivation `items` is false, so there are none).
          let prefixLen = length prefixSchemas
          prefixVals <- traverse (schemaGen defns . dereference defns) prefixSchemas
          trailing <- case schema ^. items of
            Just (OpenApiItemsObject ref) -> do
              size <- getSize
              let itemSchema = dereference defns ref
                  minLength' = fromMaybe prefixLen $ fromInteger <$> schema ^. minItems
                  maxLength' = fromMaybe (prefixLen + size) $ fromInteger <$> schema ^. maxItems
                  minTrail = max 0 (minLength' - prefixLen)
                  maxTrail = max minTrail (maxLength' - prefixLen)
              n <- choose (minTrail, maxTrail)
              vectorOf n $ schemaGen defns itemSchema
            -- items: false (or absent) -> no elements beyond the prefix.
            _ -> pure []
          return . Array . V.fromList $ prefixVals ++ trailing
      | Just items <- schema ^. items ->
          case items of
            OpenApiItemsObject ref -> do
              size <- getSize
              let itemSchema = dereference defns ref
                  minLength' = fromMaybe 0 $ fromInteger <$> schema ^. minItems
                  maxLength' = fromMaybe size $ fromInteger <$> schema ^. maxItems
              arrayLength <- choose (minLength', max minLength' maxLength')
              generatedArray <- vectorOf arrayLength $ schemaGen defns itemSchema
              return . Array $ V.fromList generatedArray
            -- items: true  -> any element allowed; items: false -> no elements.
            OpenApiItemsBoolean _ -> pure $ Array V.empty
      | otherwise -> pure $ Array V.empty
    Just OpenApiString -> do
      size <- getSize
      let minLength' = fromMaybe 0 $ fromInteger <$> schema ^. minLength
      let maxLength' = fromMaybe size $ fromInteger <$> schema ^. maxLength
      length <- choose (minLength', max minLength' maxLength')
      str <- vectorOf length arbitrary
      return . String $ T.pack str
    Just OpenApiObject -> do
      size <- getSize
      let props = dereference defns <$> schema ^. properties
          reqKeys = S.fromList $ schema ^. required
          allKeys = S.fromList . M.keys $ schema ^. properties
          optionalKeys = allKeys S.\\ reqKeys
          minProps' =
            fromMaybe (length reqKeys)
              $ fromInteger
              <$> schema
              ^. minProperties
          maxProps' = fromMaybe size $ fromInteger <$> schema ^. maxProperties
      shuffledOptional <- shuffle $ S.toList optionalKeys
      numProps <- choose (minProps', max minProps' maxProps')
      let presentKeys = take numProps $ S.toList reqKeys ++ shuffledOptional
      let presentProps = M.filterWithKey (\k _ -> k `elem` presentKeys) props
      let gens = schemaGen defns <$> presentProps
      additionalGens <- case schema ^. additionalProperties of
        Just (AdditionalPropertiesSchema addlSchema) -> do
          additionalKeys <- sequence . take (numProps - length presentProps) . repeat $ T.pack <$> arbitrary
          return . M.fromList $ zip additionalKeys (repeat . schemaGen defns $ dereference defns addlSchema)
        _ -> return []
      x <- sequence $ gens <> additionalGens
      return . Object $ fromInsOrdHashMap x

-- | Resolve a 'Referenced' value against a set of 'Definitions': return the
-- value inline if present, otherwise look the reference up by name.
--
-- /Partial:/ throws if a 'Ref' names a definition not in the map.
dereference :: Definitions a -> Referenced a -> a
dereference _ (Inline a) = a
dereference defs (Ref (Reference ref)) = fromJust $ M.lookup ref defs

-- | Generate a random JSON 'Value' conforming to the schema derived from @a@
-- via its 'ToSchema' instance. Inherits the partiality of 'schemaGen' (it may
-- 'error' when a schema's type is neither specified nor inferable).
genValue :: (ToSchema a) => Proxy a -> Gen Value
genValue p =
  let (defs, NamedSchema _ schema) = runDeclare (declareNamedSchema p) M.empty
   in schemaGen defs schema

-- | Property that asserts every value generated from @a@'s derived schema
-- parses back successfully through its 'FromJSON' instance — i.e. the schema
-- and the decoder agree.
validateFromJSON :: forall a. (ToSchema a, FromJSON a) => Proxy a -> Property
validateFromJSON p = forAll (genValue p)
  $ \val -> case parseEither parseJSON val of
    Right (_ :: a) -> succeeded
    Left err ->
      failed
        { reason = err
        }