import torch
from blinear.utils import accuracy, rmse

def test_accuracy():
    y_pred = torch.tensor([[0.1, 0.9], [0.8, 0.2]])
    y_true = torch.tensor([1, 0])

    acc = accuracy(y_pred, y_true)
    assert acc == 1.0


def test_rmse():
    y_pred = torch.tensor([1.0, 2.0])
    y_true = torch.tensor([1.0, 2.0])

    assert rmse(y_pred, y_true) == 0.0