From 1ed17ec268cf0271f20baf985f6832c50231ba75 Mon Sep 17 00:00:00 2001
From: Zain Rizvi <ZainR@meta.com>
Date: Wed, 3 May 2023 11:49:21 -0500
Subject: [PATCH] Support uploads and downloads

---
 .../actions/pytest-cache-upload/action.yml    |  3 +-
 .github/scripts/pytest_cache.py               | 52 +++++++++++++++++++
 .github/scripts/pytest_upload_cache.py        | 32 ------------
 3 files changed, 54 insertions(+), 33 deletions(-)
 create mode 100644 .github/scripts/pytest_cache.py
 delete mode 100644 .github/scripts/pytest_upload_cache.py

diff --git a/.github/actions/pytest-cache-upload/action.yml b/.github/actions/pytest-cache-upload/action.yml
index cc843cbd670ce..34966d755a63a 100644
--- a/.github/actions/pytest-cache-upload/action.yml
+++ b/.github/actions/pytest-cache-upload/action.yml
@@ -32,7 +32,8 @@ runs:
         JOB: ${{ github.job }}
         SHARD: ${{ inputs.shard }}
       run: |
-        python3 .github/scripts/pytest_upload_cache.py \
+        python3 .github/scripts/pytest_cache.py \
+          --upload \
           --cache_dir $CACHE_DIR \
           --pr_identifier $PR_IDENTIFIER \
           --workflow $WORKFLOW \
diff --git a/.github/scripts/pytest_cache.py b/.github/scripts/pytest_cache.py
new file mode 100644
index 0000000000000..67a713ceb48e4
--- /dev/null
+++ b/.github/scripts/pytest_cache.py
@@ -0,0 +1,52 @@
+import argparse
+import sys
+from pytest_caching_utils import *
+
+def main():
+    parser = argparse.ArgumentParser(description="Upload this job's the pytest cache to S3")
+    
+    mode = parser.add_mutually_exclusive_group(required=True)
+    mode.add_argument('--upload', action='store_true', help='Upload the pytest cache to S3')
+    mode.add_argument('--download', action='store_true', help='Download the pytest cache from S3, merging it with any local cache')
+
+    parser.add_argument('--cache_dir', required=True, help='Path to the folder pytest uses for its cache')
+    parser.add_argument('--pr_identifier', required=True, help='A unique PR identifier')
+    parser.add_argument('--workflow', required=True, help='The workflow name')
+    parser.add_argument('--job', required=True, help='The job name')
+    parser.add_argument('--shard', required='--upload' in sys.argv, help='The shard id') # Only required for upload
+
+    parser.add_argument('--temp_dir', required=False, help='Directory to store temp files')
+    parser.add_argument('--bucket', required=False, help='The S3 bucket to upload the cache to')
+
+    args = parser.parse_args()
+
+    if args.upload:
+        print(f"Uploading cache with args {args}")
+
+        # TODO: First check if it's even worth uploading a new cache: 
+        #    Does the cache even mark any failed tests?
+
+        upload_pytest_cache(
+            pr_identifier=args.pr_identifier, 
+            workflow=args.workflow, 
+            job=args.job, 
+            shard=args.shard, 
+            pytest_cache_dir=args.pytest_cache_dir,
+            bucket=args.bucket,
+            temp_dir=args.temp_dir,
+        )
+    
+    if args.download:
+        print(f"Downloading cache with args {args}")
+        download_pytest_cache(
+            pr_identifier=args.pr_identifier, 
+            workflow=args.workflow, 
+            job=args.job, 
+            pytest_cache_dir=args.pytest_cache_dir,
+            bucket=args.bucket,
+            temp_dir=args.temp_dir,
+        )
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
diff --git a/.github/scripts/pytest_upload_cache.py b/.github/scripts/pytest_upload_cache.py
deleted file mode 100644
index 319f17083ced7..0000000000000
--- a/.github/scripts/pytest_upload_cache.py
+++ /dev/null
@@ -1,32 +0,0 @@
-import argparse
-
-from pytest_caching_utils import *
-
-def main():
-    parser = argparse.ArgumentParser(description="Upload this job's the pytest cache to S3")
-    parser.add_argument('--cache_dir', required=True, help='Path to the folder containing the pytest cache')
-    parser.add_argument('--pr_identifier', required=True, help='A unique PR identifier')
-    parser.add_argument('--workflow', required=True, help='The workflow name')
-    parser.add_argument('--job', required=True, help='The job name')
-    parser.add_argument('--shard', required=True, help='The shard id')
-    parser.add_argument('--temp_dir', required=False, help='Directory to store temp files in')
-    parser.add_argument('--bucket', required=False, help='The S3 bucket to upload the cache to')
-    args = parser.parse_args()
-    print(args)
-
-    # TODO: First check if it's even worth uploading a new cache: 
-    #    Does the cache even mark any failed tests?
-
-    upload_pytest_cache(
-        pr_identifier=args.pr_identifier, 
-        workflow=args.workflow, 
-        job=args.job, 
-        shard=args.shard, 
-        pytest_cache_dir=args.pytest_cache_dir,
-        bucket=args.bucket,
-        temp_dir=args.temp_dir,
-    )
-
-
-if __name__ == '__main__':
-    main()
\ No newline at end of file