Arbitrary trees

arbitrary_trees.hs

import Data.Sequence (Seq((:|>)))
import Data.Tree (Tree(..))

import qualified Data.Sequence as Seq

foobar_test :: TestEnv ()
foobar_test = do
  tree :: Tree () <- liftIO . generate $ treeGen 50 `suchThat` ((>= 5) . length . flatten)
  ugids <- createUserGroupTree tree
  
  users <- foldM (\res ugid -> (res ++) <$> populateUserGroup ugid) [] ugids
  user1 <- liftIO $ generate (elements users)
  user2 <- liftIO $ generate (elements users `suchThat` (/= user1))
  
  docs <- mapM createDocument [user1, user2]
  
  chargeableItems :: [[Core.ChargeableItem]] <- liftIO
    $ replicateM 2 (generate (resize 100 arbitrary))
  traverse_
    (\(doc, cis) -> traverse_ (\ci -> chargeForItemSingle ci $ doc ^. #documentid) cis)
    (zip docs chargeableItems)

  -- ...

treeGen :: Arbitrary a => Int -> Gen (Tree a)
treeGen m = do
  t  <- arbitrary
  n  <- chooseInt (0, m `div` 2)
  ts <- vectorOf n (treeGen (m `div` 8))
  return (Node t ts)

-- | Create a `UserGroup` tree in the same shape as the given `Tree` argument.
-- Returns a list of the created `UserGroupID`s.
createUserGroupTree
  :: (MonadDB m, MonadThrow m, MonadTime m) => Tree a -> m [UserGroupID]
createUserGroupTree = fmap toList . go Nothing Seq.empty
  where
    go mParentID result (Node _ children) = do
      ugid :: UserGroupID <- view #id <$> case mParentID of
        Nothing -> Update.createUserGroup defaultUserGroup
        Just _ ->
          Update.createUserGroup (defaultChildUserGroup & #parentGroupID .~ mParentID)
      foldM (go $ Just ugid) (result :|> ugid) children

-- | Populate a `UserGroup` with some (1 - 3) randomly generated users.
populateUserGroup
  :: (CryptoRNG m, MonadDB m, MonadIO m, MonadLog m, MonadMask m, MonadTime m)
  => UserGroupID
  -> m [User]
populateUserGroup ugid = do
  n :: Positive Int <- liftIO $ generate (resize 3 arbitrary)
  replicateM (getPositive n)
             (instantiateUser $ randomUserTemplate { groupID = pure ugid })