1  import pytest
 2  
 3  from evidently.legacy.core import new_id
 4  from evidently.ui.service.storage.local.base import FSSpecBlobStorage
 5  from evidently.ui.service.storage.local.dataset import DatasetFileStorage
 6  from evidently.ui.service.type_aliases import ZERO_UUID
 7  
 8  
 9  @pytest.fixture
10  def tmp_path():
11      """Create a temporary directory."""
12      import tempfile
13  
14      with tempfile.TemporaryDirectory() as tmpdir:
15          yield tmpdir
16  
17  
18  @pytest.fixture
19  def blob_storage(tmp_path):
20      """Create blob storage."""
21      return FSSpecBlobStorage(base_path=tmp_path)
22  
23  
24  @pytest.fixture
25  def dataset_file_storage(blob_storage):
26      """Create dataset file storage."""
27      return DatasetFileStorage(dataset_blob_storage=blob_storage)
28  
29  
30  @pytest.fixture
31  def test_user_id():
32      """Create a test user ID."""
33      return ZERO_UUID
34  
35  
36  @pytest.fixture
37  def test_project_id():
38      """Create a test project ID."""
39      return new_id()
40  
41  
42  @pytest.fixture
43  def test_dataset_id():
44      """Create a test dataset ID."""
45      return new_id()
46  
47  
48  @pytest.fixture
49  def sample_data():
50      """Create sample data."""
51      return b"id,name\n1,test\n2,example"
52  
53  
54  def test_put_dataset(dataset_file_storage, test_user_id, test_project_id, test_dataset_id, sample_data):
55      """Test storing a dataset file."""
56      blob_id = dataset_file_storage.put_dataset(test_user_id, test_project_id, test_dataset_id, "test.csv", sample_data)
57      assert blob_id is not None
58      assert "datasets" in blob_id
59      assert str(test_project_id) in blob_id
60      assert str(test_dataset_id) in blob_id
61      assert "test.csv" in blob_id
62  
63  
64  def test_get_dataset(dataset_file_storage, test_user_id, test_project_id, test_dataset_id, sample_data):
65      """Test retrieving a dataset file."""
66      blob_id = dataset_file_storage.put_dataset(test_user_id, test_project_id, test_dataset_id, "test.csv", sample_data)
67      retrieved = dataset_file_storage.get_dataset(blob_id)
68      assert retrieved == sample_data
69  
70  
71  def test_check_dataset(dataset_file_storage, test_user_id, test_project_id, test_dataset_id, sample_data):
72      """Test checking if a dataset file exists."""
73      blob_id = dataset_file_storage.put_dataset(test_user_id, test_project_id, test_dataset_id, "test.csv", sample_data)
74      assert dataset_file_storage.check_dataset(blob_id) is True
75      assert dataset_file_storage.check_dataset("nonexistent/blob/id") is False
76  
77  
78  def test_remove_dataset(dataset_file_storage, test_user_id, test_project_id, test_dataset_id, sample_data):
79      """Test removing a dataset file."""
80      blob_id = dataset_file_storage.put_dataset(test_user_id, test_project_id, test_dataset_id, "test.csv", sample_data)
81      assert dataset_file_storage.check_dataset(blob_id) is True
82      dataset_file_storage.remove_dataset(blob_id)
83      assert dataset_file_storage.check_dataset(blob_id) is False
84  
85  
86  def test_get_dataset_blob_id(dataset_file_storage, test_project_id, test_dataset_id):
87      """Test getting dataset blob ID."""
88      blob_id = DatasetFileStorage.get_dataset_blob_id(test_project_id, test_dataset_id, "test_file.csv")
89      assert "datasets" in blob_id
90      assert str(test_project_id) in blob_id
91      assert str(test_dataset_id) in blob_id
92      assert "test_file.csv" in blob_id