Add artifacts_upload step.

This adds the implementation of the artifacts_upload.py step along with
unit tests. Adds the postbuild.py main function. Adds a new step to
yaml configuration file calling postbuild.py.

Tested: I ran unit tests with pytest to test for function behaviors
of git_metadata step. All test cases passed. I also ran an end-to-end
test with the sstate-cache, build was successful.

Signed-off-by: Simon-Lii <thesimonli@google.com>
Change-Id: I8d2b85bb3a07176f334f9811d0d5dd519bfd8814
diff --git a/cloudbuild.yaml b/cloudbuild.yaml
index 3489926..19b3844 100644
--- a/cloudbuild.yaml
+++ b/cloudbuild.yaml
@@ -29,18 +29,33 @@
   - name: ${_BUILDER_TOOLCHAIN_IMAGE}
     dir: ${_BUILD_SCRIPT_DIR}/${_MAIN_REPO_NAME}
     env:
-      - REPOS=${_REPOS}
       - FIRMWARE_MACHINE=${_FIRMWARE_MACHINE}
       - MAIN_REPO_NAME=${_MAIN_REPO_NAME}
       - BUILD_SCRIPT_DIR=${_BUILD_SCRIPT_DIR}
       - FIRMWARE_VERSION=${_FIRMWARE_VERSION}
       - METADATA_OBMC_REPO=${_METADATA_OBMC_REPO}
       - METADATA_REPO=${_METADATA_REPO}
+      - BUILD_SCRIPT_DIR=${_BUILD_SCRIPT_DIR}
       - GBMC_CONFIG=${_GBMC_CONFIG}
     args:
       - python3
       - /workspace/build.py
 
+  - name: ${_GCLOUD_SDK_TOOLCHAIN_IMAGE}
+    dir: ${_BUILD_SCRIPT_DIR}/${_MAIN_REPO_NAME}
+    entrypoint: /bin/bash
+    env:
+      - FIRMWARE_MACHINE=${_FIRMWARE_MACHINE}
+      - FIRMWARE_VERSION=${_FIRMWARE_VERSION}
+      - PROJECT_ID=${_PROJECT_ID}
+      - GCS_BUCKET_RELEASE=${_GCS_BUCKET_RELEASE}
+      - UPDATE_CACHE=${_UPDATE_CACHE}
+      - GCS_BUCKET=${_GCS_BUCKET}
+    args:
+      - '-c'
+      - |
+        python3 /workspace/postbuild.py
+
 timeout: 10000s
 options:
   machineType: E2_HIGHCPU_32
diff --git a/postbuild.py b/postbuild.py
new file mode 100644
index 0000000..88b3e89
--- /dev/null
+++ b/postbuild.py
@@ -0,0 +1,15 @@
+from steps import upload_artifacts
+
+class PostBuild():
+    def __init__(self):
+        self.upload_artifacts_step = upload_artifacts.UploadArtifacts()
+
+    def run_build(self):
+        self.upload_artifacts_step.run_script()
+
+def main():
+    postbuild = PostBuild()
+    postbuild.run_build()
+
+if __name__ == "__main__":
+    main()
diff --git a/steps/step.py b/steps/step.py
index f596ea9..67b6915 100644
--- a/steps/step.py
+++ b/steps/step.py
@@ -13,6 +13,8 @@
         self.metadata_obmc_repo = os.getenv("METADATA_OBMC_REPO")
         self.metadata_repo = os.getenv("METADATA_REPO")
         self.gbmc_config = os.getenv("GBMC_CONFIG")
+        self.update_cache = os.getenv("UPDATE_CACHE")
+        self.gcs_bucket_release = os.getenv("GCS_BUCKET_RELEASE")
 
     def parse_repos(self, repos_string):
         """Parses a string representation of an array of triples"""
diff --git a/steps/upload_artifacts.py b/steps/upload_artifacts.py
new file mode 100644
index 0000000..292d54b
--- /dev/null
+++ b/steps/upload_artifacts.py
@@ -0,0 +1,19 @@
+from steps import step
+import sh
+
+class UploadArtifacts(step.Step):
+    def gsutil_update_sstate_cache(self):
+        """Updates the sstate-cache with current successful build"""
+        sstate_cache_dir = f"build/{self.firmware_machine}/sstate-cache"
+        sstate_cache_bucket = f"gs://{self.gcs_bucket}/cache-unzip/{self.firmware_machine}"
+
+        sh.sh("-c", f"gsutil -mq rsync -r -d {sstate_cache_dir} {sstate_cache_bucket}")
+
+    def run_script(self):
+        if self.update_cache:
+            self.gsutil_update_sstate_cache()
+
+        artifacts_path = f"build/{self.firmware_machine}/tmp/deploy/images/{self.firmware_machine}"
+        gcs_dir = f"gs://{self.gcs_bucket_release}/{self.firmware_machine}/{self.firmware_version}"
+
+        self.gsutil("cp", artifacts_path, gcs_dir)
diff --git a/unit_tests/test_upload_artifacts.py b/unit_tests/test_upload_artifacts.py
new file mode 100644
index 0000000..0faaf66
--- /dev/null
+++ b/unit_tests/test_upload_artifacts.py
@@ -0,0 +1,82 @@
+import sh
+import unittest
+from unittest.mock import MagicMock
+from steps import upload_artifacts
+
+class TestUploadArtifacts(unittest.TestCase):
+    def setUp(self):
+        self.step = upload_artifacts.UploadArtifacts()
+        self.step.gsutil_update_sstate_cache = MagicMock()
+        self.step.gsutil = MagicMock()
+        self.init_step()
+
+    def init_step(self):
+        self.step.firmware_machine = "machine0"
+        self.step.firmware_version = "x.y.z"
+        self.step.project_id = "test-id"
+        self.step.gcs_bucket_release = "bucket-release"
+        self.step.gcs_bucket = "bucket"
+
+    def test_gsutil_with_sstate_cache(self):
+        """Tests run_script code flow with sstate_cache update enabled"""
+        self.step.update_cache = True
+        self.step.run_script()
+
+        artifacts_path = "build/machine0/tmp/deploy/images/machine0"
+        gcs_dir = "gs://bucket-release/machine0/x.y.z"
+
+        self.step.gsutil_update_sstate_cache.assert_called_with()
+        self.step.gsutil.assert_called_with("cp", artifacts_path, gcs_dir)
+
+    def test_gsutil_without_sstate_cache(self):
+        """Tests run_script code flow with sstate_cache update disabled"""
+        self.setUp()
+        self.step.update_cache = False
+        self.step.run_script()
+
+        artifacts_path = "build/machine0/tmp/deploy/images/machine0"
+        gcs_dir = "gs://bucket-release/machine0/x.y.z"
+
+        self.step.gsutil_update_sstate_cache.assert_not_called()
+        self.step.gsutil.assert_called_with("cp", artifacts_path, gcs_dir)
+
+    def test_sh_with_sstate_cache(self):
+        """Tests the arguments used to execute sh.sh with sstate-cache update enabled"""
+        self.step = upload_artifacts.UploadArtifacts()
+        self.init_step()
+        self.step.update_cache = True
+        sh.sh = MagicMock()
+        self.step.run_script()
+
+        actual_calls = sh.sh.call_args_list
+
+        sstate_cache_dir = "build/machine0/sstate-cache"
+        sstate_cache_bucket = "gs://bucket/cache-unzip/machine0"
+        call0 = f"gsutil -mq rsync -r -d {sstate_cache_dir} {sstate_cache_bucket}"
+
+        artifacts_path = "build/machine0/tmp/deploy/images/machine0"
+        gcs_dir = "gs://bucket-release/machine0/x.y.z"
+        call1 = f"gsutil -mq cp -r {artifacts_path} {gcs_dir}"
+
+        expected_calls = [(("-c", call0),), (("-c", call1),)]
+
+        self.assertEqual(actual_calls, expected_calls)
+
+    def test_sh_without_sstate_cache(self):
+        """Tests the arguments used to execute sh.sh with sstate-cache update disabled"""
+        self.step = upload_artifacts.UploadArtifacts()
+        self.init_step()
+        self.step.update_cache = False
+        sh.sh = MagicMock()
+        self.step.run_script()
+
+        artifacts_path = "build/machine0/tmp/deploy/images/machine0"
+        gcs_dir = "gs://bucket-release/machine0/x.y.z"
+
+        actual_calls = sh.sh.call_args_list
+
+        call0 = f"gsutil -mq cp -r {artifacts_path} {gcs_dir}"
+
+        expected_calls = [(("-c", call0),)]
+
+        self.assertEqual(actual_calls, expected_calls)