import torch

from data.schema import FACT_TYPE
from eval.benchmark_extractor import hard_negative_recall_at_one, recall_at_one


def test_retrieval_metrics_distinguish_hard_negative():
    labels = torch.tensor(
        [
            [FACT_TYPE, 0, 0, 4, -1, -1],
            [FACT_TYPE, 0, 0, 5, -1, -1],
        ]
    )
    candidates = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
    queries = torch.tensor([[0.9, 0.1], [0.1, 0.9]])
    assert recall_at_one(queries, labels, candidates, labels) == 1.0
    assert hard_negative_recall_at_one(queries, labels, candidates, labels) == 1.0