github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/python/tests/unit/sdk/dataset/test_label_attribute.py (about)

     1  import unittest
     2  from unittest.mock import MagicMock
     3  from aistore.sdk.dataset.label_attribute import LabelAttribute
     4  
     5  
     6  class TestLabelAttribute(unittest.TestCase):
     7      def setUp(self):
     8          self.label_identifier = MagicMock()
     9          self.label_attribute = LabelAttribute(
    10              name="label", label_identifier=self.label_identifier
    11          )
    12  
    13      def test_get_data_for_entry_with_label(self):
    14          filename = "sample_file"
    15          expected_label = "SampleLabel"
    16          self.label_identifier.return_value = expected_label
    17  
    18          key, data = self.label_attribute.get_data_for_entry(filename)
    19  
    20          expected_key = self.label_attribute.name
    21          self.assertEqual(key, expected_key)
    22          self.assertEqual(data, expected_label)
    23          self.label_identifier.assert_called_once_with(filename)
    24  
    25      def test_get_data_for_entry_no_label(self):
    26          filename = "unknown_file"
    27          self.label_identifier.return_value = None
    28  
    29          key, data = self.label_attribute.get_data_for_entry(filename)
    30  
    31          expected_key = self.label_attribute.name
    32          self.assertEqual(key, expected_key)
    33          self.assertIsNone(data)
    34          self.label_identifier.assert_called_once_with(filename)