CU-862hyd5wx Unify rosalind/vocab downloading in tests, identify and …

…fail meaningfully in case of 503
CogStack · Oct 12, 2023 · 7950c57 · 7950c57
1 parent b17be8f
commit 7950c57
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 22 deletions.
diff --git a/tests/archive_tests/test_ner_archive.py b/tests/archive_tests/test_ner_archive.py
@@ -15,6 +15,8 @@
 from medcat.linking.context_based_linker import Linker
 from medcat.config import Config
 
+from ..helper import VocabDownloader
+
 
 class NerArchiveTests(unittest.TestCase):
 
@@ -35,12 +37,9 @@ def setUp(self) -> None:
         # Check
         #assert cdb.cui2names == {'S-229004': {'movar', 'movarvirus', 'movarviruses'}, 'S-229005': {'cdb'}}
 
-        self.vocab_path = "./tmp_vocab.dat"
-        if not os.path.exists(self.vocab_path):
-            import requests
-            tmp = requests.get("https://medcat.rosalind.kcl.ac.uk/media/vocab.dat")
-            with open(self.vocab_path, 'wb') as f:
-                f.write(tmp.content)
+        downloader = VocabDownloader()
+        self.vocab_path = downloader.vocab_path
+        downloader.check_or_download()
 
         vocab = Vocab.load(self.vocab_path)
         # Make the pipeline

diff --git a/tests/medmentions/make_cdb.py b/tests/medmentions/make_cdb.py
@@ -5,6 +5,9 @@
 import logging
 import os
 
+from ..helper import VocabDownloader
+
+
 config = Config()
 config.general['log_level'] = logging.INFO
 config.general['spacy_model'] = 'en_core_sci_lg'
@@ -21,12 +24,9 @@
 from medcat.cdb import CDB
 from medcat.cat import CAT
 
-vocab_path = "./tmp_vocab.dat"
-if not os.path.exists(vocab_path):
-    import requests
-    tmp = requests.get("https://s3-eu-west-1.amazonaws.com/zkcl/vocab.dat")
-    with open(vocab_path, 'wb') as f:
-        f.write(tmp.content)
+downloader = VocabDownloader()
+vocab_path = downloader.vocab_path
+downloader.check_or_download()
 
 config = Config()
 cdb = CDB.load("./tmp_cdb.dat", config=config)

diff --git a/tests/test_ner.py b/tests/test_ner.py
@@ -14,6 +14,8 @@
 from medcat.config import Config
 from medcat.cdb import CDB
 
+from .helper import VocabDownloader
+
 
 class A_NERTests(unittest.TestCase):
     @classmethod
@@ -25,11 +27,9 @@ def setUpClass(cls):
         cls.cdb = CDB(config=cls.config)
 
         print("Set up Vocab")
-        vocab_path = "./tmp_vocab.dat"
-        if not os.path.exists(vocab_path):
-            tmp = requests.get("https://medcat.rosalind.kcl.ac.uk/media/vocab.dat")
-            with open(vocab_path, 'wb') as f:
-                f.write(tmp.content)
+        downloader = VocabDownloader()
+        vocab_path = downloader.vocab_path
+        downloader.check_or_download()
 
         cls.vocab = Vocab.load(vocab_path)
 

diff --git a/tests/test_pipe.py b/tests/test_pipe.py
@@ -17,6 +17,8 @@
 from transformers import AutoTokenizer
 
 
+from .helper import VocabDownloader
+
 
 class PipeTests(unittest.TestCase):
 
@@ -30,11 +32,9 @@ def setUpClass(cls) -> None:
         cls.config.linking['disamb_length_limit'] = 2
         cls.cdb = CDB(config=cls.config)
 
-        vocab_path = "./tmp_vocab.dat"
-        if not os.path.exists(vocab_path):
-            tmp = requests.get("https://medcat.rosalind.kcl.ac.uk/media/vocab.dat")
-            with open(vocab_path, 'wb') as f:
-                f.write(tmp.content)
+        downloader = VocabDownloader()
+        vocab_path = downloader.vocab_path
+        downloader.check_or_download()
 
         cls.vocab = Vocab.load(vocab_path)
         cls.spell_checker = BasicSpellChecker(cdb_vocab=cls.cdb.vocab, config=cls.config, data_vocab=cls.vocab)