80 lines
2.6 KiB
Haskell
80 lines
2.6 KiB
Haskell
-- Copyright 2016 TensorFlow authors.
|
|
--
|
|
-- Licensed under the Apache License, Version 2.0 (the "License");
|
|
-- you may not use this file except in compliance with the License.
|
|
-- You may obtain a copy of the License at
|
|
--
|
|
-- http://www.apache.org/licenses/LICENSE-2.0
|
|
--
|
|
-- Unless required by applicable law or agreed to in writing, software
|
|
-- distributed under the License is distributed on an "AS IS" BASIS,
|
|
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
-- See the License for the specific language governing permissions and
|
|
-- limitations under the License.
|
|
|
|
{-# LANGUAGE OverloadedStrings #-}
|
|
{-# LANGUAGE ScopedTypeVariables #-}
|
|
|
|
module Main where
|
|
|
|
import Control.Monad.IO.Class (liftIO)
|
|
import Data.Int (Int64)
|
|
import Google.Test (googleTest)
|
|
import TensorFlow.Nodes (Scalar(..))
|
|
import TensorFlow.Ops (scalar)
|
|
import TensorFlow.Queue
|
|
import TensorFlow.Session
|
|
( asyncProdNodes
|
|
, build
|
|
, buildAnd
|
|
, run
|
|
, runSession
|
|
, run_
|
|
)
|
|
import Test.Framework.Providers.HUnit (testCase)
|
|
import Test.HUnit ((@=?))
|
|
import qualified Data.ByteString as BS
|
|
|
|
-- | Test basic queue behaviors.
|
|
testBasic = testCase "testBasic" $ runSession $ do
|
|
(q :: Queue2 Int64 BS.ByteString) <- build $ makeQueue2 1 ""
|
|
buildAnd run_ (enqueue q 42 (scalar "Hi"))
|
|
x <- buildAnd run (dequeue q)
|
|
liftIO $ (Scalar 42, Scalar "Hi") @=? x
|
|
|
|
buildAnd run_ (enqueue q 56 (scalar "Bar"))
|
|
y <- buildAnd run (dequeue q)
|
|
liftIO $ (Scalar 56, Scalar "Bar") @=? y
|
|
|
|
-- | Test queue pumping.
|
|
testPump = testCase "testPump" $ runSession $ do
|
|
(deq, pump) <- build $ do
|
|
q :: Queue2 Int64 BS.ByteString <- makeQueue2 2 "ThePumpQueue"
|
|
(,) <$> dequeue q
|
|
<*> enqueue q 31 (scalar "Baz")
|
|
-- This is a realistic use. The pump inputs are pre-bound to some
|
|
-- nodes that produce values when pumped (e.g. read from a
|
|
-- file).
|
|
run_ (pump, pump)
|
|
|
|
(x, y) <- run (deq, deq)
|
|
liftIO $ (Scalar 31, Scalar "Baz") @=? x
|
|
liftIO $ (Scalar 31, Scalar "Baz") @=? y
|
|
|
|
testAsync = testCase "testAsync" $ runSession $ do
|
|
(deq, pump) <- build $ do
|
|
q :: Queue2 Int64 BS.ByteString <- makeQueue2 2 ""
|
|
(,) <$> dequeue q
|
|
<*> enqueue q 10 (scalar "Async")
|
|
-- Pumps the queue until canceled by runSession exiting.
|
|
asyncProdNodes pump
|
|
-- Picks up a couple values and verifies they are as expected.
|
|
run deq >>= liftIO . ((Scalar 10, Scalar "Async") @=?)
|
|
run deq >>= liftIO . ((Scalar 10, Scalar "Async") @=?)
|
|
|
|
main :: IO ()
|
|
main = googleTest [ testBasic
|
|
, testPump
|
|
, testAsync
|
|
]
|