@@ -10,6 +10,7 @@ module Main
1010
1111-------------------------------------------------------------------------------
1212import Control.Applicative
13+import Control.Concurrent.STM
1314import Control.Lens hiding (mapping, (.=))
1415import Control.Monad
1516import Control.Monad.IO.Class
@@ -20,6 +21,7 @@ import qualified Data.HashMap.Strict as HM
2021import qualified Data.Map as M
2122import Data.Monoid
2223import Data.Scientific
24+import Data.Time
2325import qualified Data.Vector as V
2426import Database.Bloodhound hiding (key)
2527import Network.HTTP.Client
@@ -37,30 +39,46 @@ import Katip.Scribes.ElasticSearch
3739main :: IO ()
3840main = defaultMain $ testGroup "katip-elasticsearch"
3941 [
40− withResource setupSearch teardownSearch esTests
42+ esTests
4143 , typeAnnotatedTests
4244 ]
4345
4446
4547-------------------------------------------------------------------------------
46−setupSearch :: IO (Scribe, IO ())
47−setupSearch = do
48+setupSearch :: (EsScribeCfg -> EsScribeCfg) -> IO (Scribe, IO ())
49+setupSearch modScribeCfg = do
4850 bh dropESSchema
49− mkEsScribe defaultEsScribeCfg { essAnnotateTypes = True, essIndexSettings = ixs } svr ixn mn DebugS V3
51+ mkEsScribe cfg svr ixn mn DebugS V3
52+ where
53+ cfg = modScribeCfg (defaultEsScribeCfg { essAnnotateTypes = True
54+ , essIndexSettings = ixs
55+ })
5056
5157
5258-------------------------------------------------------------------------------
5359teardownSearch :: (Scribe, IO ()) -> IO ()
5460teardownSearch (_, finalizer) = do
5561 finalizer
56− bh dropESSchema
62+ bh $ do
63+ dropESSchema
64+ dropESSTemplate
65+
66+
67+-------------------------------------------------------------------------------
68+withSearch :: (IO (Scribe, IO ()) -> TestTree) -> TestTree
69+withSearch = withSearch' id
70+
71+
72+-------------------------------------------------------------------------------
73+withSearch' :: (EsScribeCfg -> EsScribeCfg) -> (IO (Scribe, IO ()) -> TestTree) -> TestTree
74+withSearch' modScribeCfg = withResource (setupSearch modScribeCfg) teardownSearch
5775
5876
5977-------------------------------------------------------------------------------
60−esTests :: IO (Scribe, IO ()) -> TestTree
61−esTests setup = testGroup "elasticsearch scribe"
78+esTests :: TestTree
79+esTests = testGroup "elasticsearch scribe"
6280 [
63− testCase "it flushes to elasticsearch" $ withTestLogging setup $ \done -> do
81+ withSearch $ \setup -> testCase "it flushes to elasticsearch" $ withTestLogging setup $ \done -> do
6482 $(logT) (ExampleCtx True) mempty InfoS "A test message"
6583 liftIO $ do
6684 void done
@@ -69,9 +87,42 @@ esTests setup = testGroup "elasticsearch scribe"
6987 let l = head logs
7088 l ^? key "_source" . key "msg" . _String @?= Just "A test message"
7189 l ^? key "_source" . key "data" . key "whatever::b" . _Bool @?= Just True
90+ , withSearch' (\c -> c { essIndexSharding = DailyIndexSharding}) $ \setup -> testCase "date-based index sharding" $ do
91+ let t1 = mkTime 2016 1 2 3 4 5
92+ fakeClock <- newTVarIO t1
93+ withTestLogging' (set logEnvTimer (readTVarIO fakeClock)) setup $ \done -> do
94+ $(logT) (ExampleCtx True) mempty InfoS "today"
95+ let t2 = mkTime 2016 1 3 3 4 5
96+ liftIO (atomically (writeTVar fakeClock t2))
97+ $(logT) (ExampleCtx True) mempty InfoS "tomorrow"
98+ liftIO $ do
99+ void done
100+ todayLogs <- getLogsByIndex (IndexName "katip-elasticsearch-tests-2016-1-2")
101+ tomorrowLogs <- getLogsByIndex (IndexName "katip-elasticsearch-tests-2016-1-3")
102+ assertBool ("todayLogs has " <> show (length todayLogs) <> " items") (length todayLogs == 1)
103+ assertBool ("tomorrowLogs has " <> show (length tomorrowLogs) <> " items") (length tomorrowLogs == 1)
104+ let logToday = head todayLogs
105+ let logTomorrow = head tomorrowLogs
106+ logToday ^? key "_source" . key "msg" . _String @?= Just "today"
107+ logTomorrow ^? key "_source" . key "msg" . _String @?= Just "tomorrow"
72108 ]
73109
74110
111+-------------------------------------------------------------------------------
112+mkTime :: Integer -> Int -> Int -> DiffTime -> DiffTime -> DiffTime -> UTCTime
113+mkTime y m d hr minute s = UTCTime day dt
114+ where
115+ day = mkDay y m d
116+ dt = s + 60 * minute + 60 * 60 * hr
117+
118+
119+-------------------------------------------------------------------------------
120+mkDay :: Integer -> Int -> Int -> Day
121+mkDay y m d = day
122+ where
123+ Just day = fromGregorianValid y m d
124+
125+
75126-------------------------------------------------------------------------------
76127data ExampleCtx = ExampleCtx {
77128 ecBool :: Bool
@@ -133,9 +184,17 @@ annotatedExampleValue = Array $ V.fromList
133184
134185-------------------------------------------------------------------------------
135186getLogs :: IO [Value]
136−getLogs = do
137− r <- bh $ searchByIndex ixn $ mkSearch Nothing Nothing
138− statusCode (responseStatus r) @?= 200
187+getLogs = getLogsByIndex ixn
188+
189+
190+-------------------------------------------------------------------------------
191+getLogsByIndex :: IndexName -> IO [Value]
192+getLogsByIndex i = do
193+ r <- bh $ do
194+ void (refreshIndex i)
195+ searchByIndex i (mkSearch Nothing Nothing)
196+ let actualCode = statusCode (responseStatus r)
197+ assertBool ("search by " <> show i <> " " <> show actualCode <> " /= 200") (actualCode == 200)
139198 return $ responseBody r ^.. key "hits" . key "hits" . values
140199
141200
@@ -147,11 +206,20 @@ bh = withBH defaultManagerSettings svr
147206-------------------------------------------------------------------------------
148207withTestLogging
149208 :: IO (Scribe, IO a) -> (IO Reply -> KatipT IO b) -> IO b
150−withTestLogging setup f = do
151− (scr, done) <- setup
152− le <- initLogEnv ns env
153− let done' = done >> bh (refreshIndex ixn)
154− runKatipT le { _logEnvScribes = M.singleton "es" scr} (f done')
209+withTestLogging = withTestLogging' id
210+
211+
212+-------------------------------------------------------------------------------
213+withTestLogging'
214+ :: (LogEnv -> LogEnv)
215+ -> IO (Scribe, IO a)
216+ -> (IO Reply -> KatipT IO b)
217+ -> IO b
218+withTestLogging' modEnv setup f = do
219+ (scr, done) <- setup
220+ le <- modEnv <$> initLogEnv ns env
221+ let done' = done >> bh (refreshIndex ixn)
222+ runKatipT le { _logEnvScribes = M.singleton "es" scr} (f done')
155223 where
156224 ns = Namespace ["katip-test"]
157225 env = Environment "test"
@@ -179,7 +247,12 @@ mn = MappingName "logs"
179247
180248-------------------------------------------------------------------------------
181249dropESSchema :: BH IO ()
182−dropESSchema = void $ deleteIndex ixn
250+dropESSchema = void $ deleteIndex (IndexName "katip-elasticsearch-tests*")
251+
252+
253+-------------------------------------------------------------------------------
254+dropESSTemplate :: BH IO ()
255+dropESSTemplate = void $ deleteTemplate (TemplateName "katip-elasticsearch-tests")
183256
184257
185258-------------------------------------------------------------------------------