Helpers.hs

{-# LANGUAGE QuasiQuotes #-}

module Test.Integration.PGMQ.Helpers
  ( pgmqHelpersTests )
where

import Control.Exception (bracket)
import Data.Text qualified as T
import Database.PostgreSQL.Simple qualified as PSQL
import Database.PostgreSQL.Simple.Newtypes qualified as PSQL (Aeson(..))
import Database.PostgreSQL.Simple.SqlQQ (sql)
import Database.PostgreSQL.Simple.Types qualified as PSQL (QualifiedIdentifier(..))
import Database.PGMQ qualified as PGMQ
import Test.Hspec
import Test.Integration.Utils (getPSQLEnvConnectInfo, randomQueueName)
import Test.RandomStrings (randomASCII, randomString, onlyAlphaNum)


data TestEnv =
  TestEnv {
    conn  :: PSQL.Connection
  , queue :: PGMQ.Queue
  }
    
-- NOTE These tests expect a local pgmq server runnign on port 5432.

testQueuePrefix :: PGMQ.Queue
testQueuePrefix = "test_pgmq"

setUpConn :: IO TestEnv
setUpConn = do
  connInfo <- getPSQLEnvConnectInfo
  conn <- PSQL.connect connInfo
  queue <- randomQueueName testQueuePrefix
  return $ TestEnv { conn, queue }

dropConn :: TestEnv -> IO ()
dropConn (TestEnv { conn }) = do
  PSQL.close conn

withConn :: (TestEnv -> IO ()) -> IO ()
withConn = bracket setUpConn dropConn

withPGMQ :: (TestEnv -> IO ()) -> IO ()
withPGMQ f = withConn $ \testEnv -> bracket (setUpPGMQ testEnv) (tearDownPGMQ testEnv) (\_ -> f testEnv)
  where
    setUpPGMQ (TestEnv { conn, queue }) = do
      PGMQ.initialize conn
      PGMQ.createQueue conn queue

    tearDownPGMQ (TestEnv { conn = _conn, queue = _queue }) _ = do
      -- PGMQ.dropQueue conn queue
      pure ()

      
pgmqHelpersTests :: Spec
pgmqHelpersTests = parallel $ around withPGMQ $ describe "PGMQ Helpers" $ do
  it "can send message when queue doesn't exist" $ \(TestEnv { conn, queue }) -> do
    let newQueue = queue <> "_new"
    let queueId = PSQL.QualifiedIdentifier (Just "pgmq") $ T.pack ("q_" <> newQueue)
    let iter = [1..20] :: [Int]  -- number of steps
    mapM_ (\_i -> do
              -- Generate random strings and make sure that the
              -- message ids we get from sendMessage match our data
              message <- randomString (onlyAlphaNum randomASCII) 20
              msgId <- PGMQ.safeSendMessage conn newQueue message 0
              [PSQL.Only (PSQL.Aeson msg)] <- PSQL.query conn [sql| SELECT message FROM ? WHERE msg_id = ? |]
                                                              (queueId, msgId)
              msg `shouldBe` message
              ) iter