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

     1  import unittest
     2  from unittest.mock import patch
     3  from pathlib import Path
     4  from aistore.sdk.dataset.data_attribute import DataAttribute
     5  
     6  
     7  class TestDataAttribute(unittest.TestCase):
     8      def setUp(self):
     9          self.data_attribute = DataAttribute(
    10              path=Path("/fake/path"), name="testfile", file_type="jpg"
    11          )
    12  
    13      @patch("aistore.sdk.dataset.data_attribute.read_file_bytes")
    14      def test_get_data_for_entry_file_not_found(self, mock_read_file_bytes):
    15          filename = "nonexistent"
    16          expected_key = f"{self.data_attribute.name}.{self.data_attribute.file_type}"
    17          mock_read_file_bytes.side_effect = FileNotFoundError("No such file")
    18          key, data = self.data_attribute.get_data_for_entry(filename)
    19  
    20          self.assertEqual(key, expected_key)
    21          self.assertIsNone(data)
    22          mock_read_file_bytes.assert_called_once_with(Path("/fake/path/nonexistent.jpg"))
    23  
    24      @patch("aistore.sdk.dataset.data_attribute.read_file_bytes")
    25      def test_get_data_for_entry_success(self, mock_read_file_bytes):
    26          filename = "sample"
    27          expected_key = f"{self.data_attribute.name}.{self.data_attribute.file_type}"
    28          expected_data = b"fake data"
    29          mock_read_file_bytes.return_value = expected_data
    30          key, data = self.data_attribute.get_data_for_entry(filename)
    31  
    32          self.assertEqual(key, expected_key)
    33          self.assertEqual(data, expected_data)
    34          mock_read_file_bytes.assert_called_once_with(Path("/fake/path/sample.jpg"))