wip mock ffs

Author: matt-bernstein

label_studio/io_storages/tests/test_multitask_import.py

@@ -1,6 +1,7 @@
 import json
 
 import boto3
+import mock
 import pytest
 from io_storages.models import S3ImportStorage
 from io_storages.s3.models import S3ImportStorageLink
@@ -14,13 +15,55 @@
 from moto import mock_s3
 from projects.tests.factories import ProjectFactory
 from rest_framework.test import APIClient
-from tests.utils import azure_client_mock, gcs_client_mock, mock_feature_flag, redis_client_mock
+
+from tests.utils import azure_client_mock, gcs_client_mock, redis_client_mock
+
+# FIXME: mock_feature_flag is not working in this file for some reason
+
+
+@pytest.fixture(name='fflag_feat_dia_2092_multitasks_per_storage_link_on')
+def fflag_feat_dia_2092_multitasks_per_storage_link_on():
+    from core.feature_flags import flag_set
+
+    def fake_flag_set(*args, **kwargs):
+        if args[0] == 'fflag_feat_dia_2092_multitasks_per_storage_link':
+            return True
+        return flag_set(*args, **kwargs)
+
+    with mock.patch('core.feature_flags.flag_set', wraps=fake_flag_set):
+        yield
+
+
+@pytest.fixture(name='fflag_feat_root_11_support_jsonl_cloud_storage_on')
+def fflag_feat_root_11_support_jsonl_cloud_storage_on():
+    from core.feature_flags import flag_set
+
+    def fake_flag_set(*args, **kwargs):
+        if args[0] == 'fflag_feat_root_11_support_jsonl_cloud_storage':
+            return True
+        return flag_set(*args, **kwargs)
+
+    with mock.patch('core.feature_flags.flag_set', wraps=fake_flag_set):
+        yield
+
+
+@pytest.fixture(name='fflag_feat_root_11_support_jsonl_cloud_storage_off')
+def fflag_feat_root_11_support_jsonl_cloud_storage_off():
+    from core.feature_flags import flag_set
+
+    def fake_flag_set(*args, **kwargs):
+        if args[0] == 'fflag_feat_root_11_support_jsonl_cloud_storage':
+            return False
+        return flag_set(*args, **kwargs)
+
+    with mock.patch('core.feature_flags.flag_set', wraps=fake_flag_set):
+        yield
+
 
 #
 # Integration tests for storage.sync()
 #
 
-
 pytestmark = pytest.mark.django_db
 
 
@@ -66,7 +109,7 @@ def _test_storage_import(project, storage_class, task_data, **storage_kwargs):
         assert task['data'] == expected_data['data']
 
 
-@mock_feature_flag('fflag_feat_dia_2092_multitasks_per_storage_link', True)
+@pytest.mark.fflag_feat_dia_2092_multitasks_per_storage_link_on
 def test_import_multiple_tasks_s3(project, common_task_data):
     with mock_s3():
         # Setup S3 bucket and test data
@@ -88,7 +131,7 @@ def test_import_multiple_tasks_s3(project, common_task_data):
         )
 
 
-@mock_feature_flag('fflag_feat_dia_2092_multitasks_per_storage_link', True)
+@pytest.mark.fflag_feat_dia_2092_multitasks_per_storage_link_on
 def test_import_multiple_tasks_gcs(project, common_task_data):
     # initialize mock with sample data
     with gcs_client_mock():
@@ -102,7 +145,7 @@ def test_import_multiple_tasks_gcs(project, common_task_data):
         )
 
 
-@mock_feature_flag('fflag_feat_dia_2092_multitasks_per_storage_link', True)
+@pytest.mark.fflag_feat_dia_2092_multitasks_per_storage_link_on
 def test_import_multiple_tasks_azure(project, common_task_data):
     # initialize mock with sample data
     with azure_client_mock(sample_json_contents=common_task_data, sample_blob_names=['test.json']):
@@ -114,7 +157,7 @@ def test_import_multiple_tasks_azure(project, common_task_data):
         )
 
 
-@mock_feature_flag('fflag_feat_dia_2092_multitasks_per_storage_link', True)
+@pytest.mark.fflag_feat_dia_2092_multitasks_per_storage_link_on
 def test_import_multiple_tasks_redis(project, common_task_data):
     with redis_client_mock() as redis:
         redis.set('test.json', json.dumps(common_task_data))
@@ -128,7 +171,7 @@ def test_import_multiple_tasks_redis(project, common_task_data):
         )
 
 
-@mock_feature_flag('fflag_feat_dia_2092_multitasks_per_storage_link', True)
+@pytest.mark.fflag_feat_dia_2092_multitasks_per_storage_link_on
 def test_storagelink_fields(project, common_task_data):
     # use an actual storage and storagelink to test this, since factories aren't connected properly
     with mock_s3():
@@ -313,7 +356,7 @@ def test_list_jsonl(storage):
     create_tasks(storage, output)
 
 
-@mock_feature_flag('fflag_feat_root_11_support_jsonl_cloud_storage', True)
+@pytest.mark.fflag_feat_root_11_support_jsonl_cloud_storage_on
 def test_list_jsonl_with_preds_and_annots(storage):
     task_data = annots_preds_task_list
 
@@ -331,7 +374,7 @@ def test_list_jsonl_with_preds_and_annots(storage):
     create_tasks(storage, output)
 
 
-@mock_feature_flag('fflag_feat_root_11_support_jsonl_cloud_storage', False)
+@pytest.mark.fflag_feat_root_11_support_jsonl_cloud_storage_off
 def test_ff_blocks_jsonl():
     with pytest.raises(ValueError):
         load_tasks_json(b'{"text": "Test task 1"}\n{"text": "Test task 2"}', 'test.jsonl')