import unittest

from metrics.geometric_mean.geometric_mean import GeometricMean

geometric_mean = GeometricMean()


class TestGeometricMean(unittest.TestCase):
    def test_gmean(self):
        refs = [0, 1, 2, 0, 1, 2]
        preds = [0, 1, 2, 0, 1, 2]
        geometric_mean_score = geometric_mean.compute(predictions=preds, references=refs)
        print(geometric_mean_score)
        self.assertTrue(geometric_mean_score == {'geometric-mean': 1.0})

        refs = [0, 2, 1, 0, 0, 1]
        preds = [0, 1, 2, 0, 1, 2]
        geometric_mean_score = geometric_mean.compute(predictions=preds, references=refs)
        self.assertTrue(geometric_mean_score == {'geometric-mean': 0.0})