haskell-opentimestamps-0.5.5.0: test/Arbitrary.hs

{-# LANGUAGE GADTs #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Arbitrary
  ( Arbitrary (..),
    arbitrarySizedByteString
  ) where

import Bitcoin.Core.RPC (BlockHeader (..))
import qualified Data.ByteString as BS
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import Data.Time.Clock.POSIX (posixSecondsToUTCTime)
import Data.Word (Word32)
import Haskoin.Block.Common (BlockHash (..))
import Haskoin.Crypto.Hash (sha256)
import OpenTimestamps.Attestation (Attestation (..))
import OpenTimestamps.DetachedTimestampFile
  ( DetachedTimestampFile (..)
  , DigestType (..)
  )
import OpenTimestamps.Op
import OpenTimestamps.Timestamp (Timestamp (..))
import Test.QuickCheck
  ( Arbitrary (..)
  , Gen
  , choose
  , elements
  , listOf
  , oneof
  , resize
  , sized
  , vectorOf
  )
import qualified Data.Text as T

instance Arbitrary Op where
  arbitrary =
    oneof
      [ pure Sha1
      , pure Sha256
      , pure Ripemd160
      , pure Keccak256
      , pure Hexlify
      , pure Reverse
      , Append <$> arbitraryBS
      , Prepend <$> arbitraryBS
      ]
    where
      arbitraryBS :: Gen BS.ByteString
      arbitraryBS = BS.pack <$> vectorOf 10 arbitrary

instance Arbitrary Attestation where
  arbitrary =
    oneof
      [ Bitcoin <$> arbitrary
      , Pending . T.pack <$> listOf (elements ['a'..'z'])
      , Unknown <$> arbitrarySizedByteString (1, 8) <*> arbitrarySizedByteString (0, 32)
      ]

instance Arbitrary Timestamp where
  arbitrary = sized arbitraryTimestamp
    where
      arbitraryTimestamp :: Int -> Gen Timestamp
      arbitraryTimestamp n
        | n <= 0 =
            Timestamp
              <$> arbitrarySizedByteString (1, 32)
              <*> (Set.fromList <$> listOf arbitrary)
              <*> pure Map.empty
        | otherwise =
            oneof
              [ Timestamp
                  <$> arbitrarySizedByteString (1, 32)
                  <*> (Set.fromList <$> listOf arbitrary)
                  <*> pure Map.empty
              , Timestamp
                  <$> arbitrarySizedByteString (1, 32)
                  <*> (Set.fromList <$> listOf arbitrary)
                  <*> (Map.fromList <$> listOf ((,) <$> arbitrary <*> resize (n `div` 2) (arbitraryTimestamp (n `div` 2))))
              ]

instance Arbitrary DetachedTimestampFile where
  arbitrary = do
    digestType' <- arbitrary
    timestamp' <- arbitrary
    pure $ DetachedTimestampFile
      { digestType = digestType'
      , timestamp = timestamp'
      }

instance Arbitrary DigestType where
  arbitrary = elements [DSha256, DSha1, DRipemd160]

instance Arbitrary BlockHash where
  arbitrary = BlockHash . sha256 <$> arbitrarySizedByteString (32, 32)

instance Arbitrary BlockHeader where
  arbitrary = do
    hash <- arbitrary
    confs <- choose (0 :: Word32, 1000 :: Word32)
    height <- choose (0 :: Word32, 1000000 :: Word32)
    merkleRoot <- sha256 <$> arbitrarySizedByteString (32, 32)
    time <- posixSecondsToUTCTime . fromIntegral <$> (choose (0, 2147483647) :: Gen Word32)
    medianTime <- posixSecondsToUTCTime . fromIntegral <$> (choose (0, 2147483647) :: Gen Word32)
    nonce <- arbitrary
    difficulty <- arbitrary
    txCount <- choose (0 :: Int, 10000 :: Int)
    prevHash <- oneof [pure Nothing, Just <$> arbitrary]
    nextHash <- oneof [pure Nothing, Just <$> arbitrary]
    pure $ BlockHeader
      { blockHeaderHash = hash
      , blockHeaderConfs = confs
      , blockHeaderHeight = height
      , blockHeaderMerkleRoot = merkleRoot
      , blockHeaderTime = time
      , blockHeaderMedianTime = medianTime
      , blockHeaderNonce = nonce
      , blockHeaderDifficulty = difficulty
      , blockHeaderTxCount = txCount
      , blockHeaderPrevHash = prevHash
      , blockHeaderNextHash = nextHash
      }

instance Arbitrary BS.ByteString where
  arbitrary = arbitrarySizedByteString (1, 32)

-- | Helper to generate a ByteString of a specific size range.
arbitrarySizedByteString :: (Int, Int) -> Gen BS.ByteString
arbitrarySizedByteString (minSize, maxSize) = do
  size <- choose (minSize, maxSize)
  BS.pack <$> vectorOf size arbitrary