1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
|
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
import Data.Aeson (Value (..), eitherDecode, encode)
import Data.Aeson.Diff as AesonDiff
import qualified Data.ByteString.Base64 as Base64
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as B
import qualified Data.HashMap.Strict as HM
import Data.Ipynb
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.Vector as V
import Lens.Micro
import Lens.Micro.Aeson
import System.Directory
import System.FilePath
import Test.Tasty
import Test.Tasty.HUnit
import Data.Monoid
main :: IO ()
main = do
let rtdir = "test" </> "rt-files"
createDirectoryIfMissing False rtdir
fs <- map (rtdir </>) . filter isIpynb <$> getDirectoryContents rtdir
defaultMain $ testGroup "round-trip tests" $ map rtTest fs
isIpynb :: FilePath -> Bool
isIpynb fp = takeExtension fp == ".ipynb"
-- We don't want tests failing because of inconsequential
-- differences in formatting of base64 data, like line breaks.
normalizeBase64 :: Value -> Value
normalizeBase64 bs =
bs & key "cells" . values . key "outputs" . values . key "data"
. _Object %~ HM.mapWithKey (\k v ->
if k == "application/json" ||
"text/" `T.isPrefixOf` k ||
"+json" `T.isSuffixOf` k
then v
else go v)
where
go (Array vec) =
go $ String
$ mconcat $ map
(\v' -> case v' of
String t' -> t'
_ -> error "expected String") $ V.toList vec
go (String t) =
case Base64.decode (TE.encodeUtf8 (T.filter (/='\n') t)) of
Left _ -> String t -- textual
Right b -> String $
TE.decodeUtf8 . (<> "\n") .
B.intercalate "\n" . chunksOf 76 .
Base64.encode $ b
go v = v
chunksOf :: Int -> B.ByteString -> [B.ByteString]
chunksOf k s
| B.null s = []
| otherwise =
let (h,t) = B.splitAt k s
in h : chunksOf k t
rtTest :: FilePath -> TestTree
rtTest fp = testCase fp $ do
inRaw <- BL.readFile fp
let format = inRaw ^? key "nbformat"._Number
case format of
Just 4 -> rtTest4 inRaw
_ -> rtTest3 inRaw
rtTest3 :: BL.ByteString -> IO ()
rtTest3 inRaw = do
(inJSON :: Value) <- either error return $ eitherDecode inRaw
(nb :: Notebook NbV3) <- either error return $ eitherDecode inRaw
let outRaw = encode nb
(nb' :: Notebook NbV3) <- either error return $ eitherDecode outRaw
(outJSON :: Value) <- either error return $ eitherDecode outRaw
-- test that (read . write) == id
let patch' = AesonDiff.diff
(normalizeBase64 inJSON) (normalizeBase64 outJSON)
assertBool (show patch') (patch' == Patch [])
-- now test that (write . read) == id
assertEqual "write . read != read" nb nb'
rtTest4 :: BL.ByteString -> IO ()
rtTest4 inRaw = do
(inJSON :: Value) <- either error return $ eitherDecode inRaw
(nb :: Notebook NbV4) <- either error return $ eitherDecode inRaw
let outRaw = encode nb
(nb' :: Notebook NbV4) <- either error return $ eitherDecode outRaw
(outJSON :: Value) <- either error return $ eitherDecode outRaw
-- test that (read . write) == id
let patch' = AesonDiff.diff
(normalizeBase64 inJSON) (normalizeBase64 outJSON)
assertBool (show patch') (patch' == Patch [])
-- now test that (write . read) == id
assertEqual "write . read != read" nb nb'
|
