✅ Update tests

Author: znatty22

tests/files/test_file_models.py

@@ -2,10 +2,11 @@
 from django.contrib.auth import get_user_model
 from graphql_relay import to_global_id
 
+from creator.studies.models import Study
 from creator.studies.factories import StudyFactory
 from creator.files.models import File
 
-from creator.files.factories import FileFactory
+from creator.files.factories import FileFactory, VersionFactory
 
 User = get_user_model()
 
@@ -40,3 +41,35 @@ def test_s3_scrape(db, clients, upload_file):
 
     assert "S3S" in file.valid_types
     assert file.valid_types == version.valid_types
+
+
+def test_set_storage(db, settings):
+    """
+    Test Version.set_storage
+    """
+    # File storage
+    fv = VersionFactory()
+    fv.set_storage()
+    assert "FileSystemStorage" in str(fv.key.storage)
+
+    # S3 storage - get study from root_file
+    s3_storage = "django_s3_storage.storage.S3Storage"
+    settings.DEFAULT_FILE_STORAGE = s3_storage
+    fv.set_storage()
+    assert fv.key.storage.settings.AWS_S3_BUCKET_NAME == (
+        fv.root_file.study.bucket
+    )
+
+    # S3 storage - get study from prop
+    study = fv.root_file.study
+    fv.study = study
+    fv.set_storage()
+    assert fv.key.storage.settings.AWS_S3_BUCKET_NAME == (
+        fv.root_file.study.bucket
+    )
+
+    # No study
+    fv.root_file = None
+    fv.study = None
+    with pytest.raises(Study.DoesNotExist):
+        fv.set_storage()

tests/ingest_runs/test_validation.py

@@ -300,12 +300,16 @@ class MockVersion:
         def __init__(self):
             self.pk = 1
             self.file_name = "MyVersion"
+            self.key = f"{self.file_name}.txt"
+
+        def set_storage(self):
+            pass
 
     clean_and_map = validation_run.clean_and_map
 
     # Error extracting file content into DataFrame
     mocker.patch(
-        "creator.ingest_runs.tasks.validation_run.extract_data",
+        "creator.ingest_runs.tasks.validation_run.read_df",
         side_effect=Exception,
     )
     with pytest.raises(validation_run.ExtractDataError) as e:
@@ -320,7 +324,7 @@ def __init__(self):
         }
     )
     mock_extract_data = mocker.patch(
-        "creator.ingest_runs.tasks.validation_run.extract_data",
+        "creator.ingest_runs.tasks.validation_run.read_df",
         return_value=df
     )
     mapped_df = clean_and_map(MockVersion(), mapper)