module TonelliShanksTest (run) where

import Test.HUnit

import ModularSquareRoot (tonelliShanks)

assertOneOf :: (Eq a, Show a) => String -> Maybe a -> [a] -> Assertion
assertOneOf msg result expected =
    case result of
        Nothing -> assertFailure (msg ++ ": expected one of " ++ show expected ++ ", got Nothing")
        Just x -> assertBool
                    (msg ++ ": got " ++ show x ++ ", expected one of " ++ show expected)
                    (x `elem` expected)

run :: IO ()
run = do
    putStrLn "Tonelli Shanks test"

    assertOneOf "sqrt(1) mod 7" (tonelliShanks 1 7) [1,6]
    assertOneOf "sqrt(2) mod 7" (tonelliShanks 2 7) [3,4]
    assertEqual "sqrt(3) mod 7" (tonelliShanks 3 7) Nothing
    assertOneOf "sqrt(4) mod 7" (tonelliShanks 4 7) [2,5]
    assertEqual "sqrt(5) mod 7" (tonelliShanks 5 7) Nothing
    assertEqual "sqrt(6) mod 7" (tonelliShanks 6 7) Nothing

    putStrLn "Tonelli Shanks test done."