import remedapy as R


class TestWhen:
    def test_data_first(self):
        # when(data, predicate, onTrue, ...extraArgs)
        # when(data, predicate, { onTrue, onFalse }, ...extraArgs)
        assert R.when(4, R.gt(3), R.add(1)) == 5
        assert R.when(2, R.gt(3), R.add(1)) == 2
        assert R.when(2, R.gt(3), R.add(1), on_false=R.multiply(2)) == 4
        assert R.when(2, R.gt(3), R.add(1), on_false=R.constant(5)) == 5

    def test_data_last(self):
        # when(predicate, onTrue)(data, ...extraArgs)
        # when(predicate, { onTrue, onFalse })(data, ...extraArgs)
        x = R.pipe(None, R.when(R.is_none, R.constant(42)))
        assert x == 42
        fn = R.when(R.gt(3), R.add(1), on_false=R.multiply(2))
        x1 = fn(2)  # this should be int
        assert x1 == 4
        x2 = fn(5)
        assert x2 == 6

    def test_typing(self):
        # TODO: research if it is possible to make pyright properly infer the return type
        fn = R.when(R.gt(3), R.identity, on_false=R.identity)
        x1: int | float = fn(1)  # this should be int
        assert x1 == 1
        x2: int | float = fn(3.5)  # this should be float
        assert x2 == 3.5
