improve hf hub push test (#35)

fcakyon · Jan 22, 2023 · 0703b9e · 0703b9e
1 parent 04219a0
commit 0703b9e
Showing 1 changed file with 14 additions and 13 deletions.
diff --git a/tests/test_hf.py b/tests/test_hf.py
@@ -69,12 +69,9 @@ def run(cmd):
 def test_detection_upload():
     import platform
     from packaging.version import Version
+    from huggingface_hub.utils._errors import HfHubHTTPError
 
     # run following lines if linux and python major == 3 and python minor == 10 (python micor can be anything)
-    print(f'platform.system(): {platform.system()}')
-    print(f'platform.python_version(): {platform.python_version()}')
-    print(f'os.getenv("RUNNER_OS"): {os.getenv("RUNNER_OS")}')
-    print(f'Version(platform.python_version()) >= Version("3.10"): {Version(platform.python_version()) >= Version("3.10")}')
     if platform.system() == 'Linux' and Version(platform.python_version()) >= Version("3.10"):
         print('training started')
         run('yolo train detect model=yolov8n.pt data=coco8.yaml imgsz=32 epochs=1')
@@ -83,12 +80,16 @@ def test_detection_upload():
         if hf_token is None:
             raise ValueError('Please set HF_TOKEN environment variable to your HuggingFace token.')
         print('push to hub started')
-        push_to_hfhub(
-            hf_model_id="fcakyon/yolov8n-test",
-            exp_dir='runs/detect/train',
-            hf_token=os.getenv('HF_TOKEN'),
-            hf_private=True,
-            hf_dataset_id="fcakyon/football-detection",
-            thumbnail_text='YOLOv8s Football Detection'
-        )
-        print('push to hub ended')
+        try:
+            push_to_hfhub(
+                hf_model_id="fcakyon/yolov8n-test",
+                exp_dir='runs/detect/train',
+                hf_token=os.getenv('HF_TOKEN'),
+                hf_private=True,
+                hf_dataset_id="fcakyon/football-detection",
+                thumbnail_text='YOLOv8s Football Detection'
+            )
+            print('push to hub succeeded')
+        except HfHubHTTPError as e:
+            print('push to hub failed')
+            print(e)