Refactor machine_setup step and add gsutil_setup step.

This refactors machine_setup to be test driven.
Adds the gsutil_setup step that creates the sstate-cache and the
sources mirror URL directories.
Adds support to decide between sstate-cache or clean build.
Adds unit tests for machine_setup and gsutil_setup.

Tested: I ran unit tests with pytest to test for function behaviors
of all the prebuild classes. All test cases passed. Performed manual
end-to-end tests on GCB with sstate-cache, all test cases passed.

Signed-off-by: Simon-Lii <thesimonli@google.com>
Change-Id: I392e100119e002af6b2c49a9e074066eb1bbfe5b
diff --git a/cloudbuild.yaml b/cloudbuild.yaml
index 9c560b0..842a652 100644
--- a/cloudbuild.yaml
+++ b/cloudbuild.yaml
@@ -1,17 +1,30 @@
 steps:
   - name: ${_GIT_CLONE_TOOLCHAIN_IMAGE}
-    entrypoint: /bin/python3
+    entrypoint: /bin/bash
+    dir: ${_BUILD_SCRIPT_DIR}
     env:
       - REPOS=${_REPOS}
       - FIRMWARE_MACHINE=${_FIRMWARE_MACHINE}
       - MAIN_REPO_NAME=${_MAIN_REPO_NAME}
       - BUILD_SCRIPT_DIR=${_BUILD_SCRIPT_DIR}
+      - CURRENT_TOOLCHAIN_IMAGE=GIT
     args:
       - -c
       - |
-        import prebuild
+        python3 /workspace/prebuild.py
 
-        prebuild.main()
+  - name: ${_GCLOUD_SDK_TOOLCHAIN_IMAGE}
+    entrypoint: /bin/bash
+    dir: ${_BUILD_SCRIPT_DIR}/${_MAIN_REPO_NAME}
+    env:
+      - FIRMWARE_MACHINE=${_FIRMWARE_MACHINE}
+      - USE_CACHE=${_USE_CACHE}
+      - CURRENT_TOOLCHAIN_IMAGE=GCLOUD
+      - GCS_BUCKET=${_GCS_BUCKET}
+    args:
+      - -c
+      - |
+        python3 /workspace/prebuild.py
 
 timeout: 10000s
 options:
diff --git a/prebuild.py b/prebuild.py
index c949266..3bfac8a 100644
--- a/prebuild.py
+++ b/prebuild.py
@@ -1,13 +1,26 @@
-from steps import git_metadata, machine_setup
+from steps import git_metadata, machine_setup, gsutil_setup
+import os
+from enum import Enum
+
+class CurrentToolchain(Enum):
+    GIT = "GIT"
+    GCLOUD = "GCLOUD"
 
 class Prebuild():
     def __init__(self):
-        self.git_metadata_step = git_metadata.GitMetadata()
-        self.machine_setup_step = machine_setup.MachineSetup()
+        self.current_toolchain_image = os.getenv("CURRENT_TOOLCHAIN_IMAGE")
 
     def run_prebuild(self):
-        self.git_metadata_step.run_script()
-        self.machine_setup_step.run_script()
+        if CurrentToolchain(self.current_toolchain_image) == CurrentToolchain.GIT:
+            git_metadata_step = git_metadata.GitMetadata()
+            machine_setup_step = machine_setup.MachineSetup()
+
+            git_metadata_step.run_script()
+            machine_setup_step.run_script()
+
+        elif CurrentToolchain(self.current_toolchain_image) == CurrentToolchain.GCLOUD:
+            gsutil_setup_step = gsutil_setup.GsutilSetup()
+            gsutil_setup_step.run_script()
 
 def main():
     prebuild = Prebuild()
diff --git a/steps/git_metadata.py b/steps/git_metadata.py
index 54b7518..6fc2155 100644
--- a/steps/git_metadata.py
+++ b/steps/git_metadata.py
@@ -11,7 +11,6 @@
         sh.git("-C", dir, "clone", repo_url)
 
     def run_script(self):
-        sh.cd(self.build_script_dir)
         for repo in self.parse_repos(self.repos):
             repo_url, branch, dir = repo
             if branch:
diff --git a/steps/gsutil_setup.py b/steps/gsutil_setup.py
new file mode 100644
index 0000000..fc10b75
--- /dev/null
+++ b/steps/gsutil_setup.py
@@ -0,0 +1,40 @@
+from steps import step
+import sh
+
+class GsutilSetup(step.Step):
+    def create_sstate_cache_dir(self):
+        """Creates the directory sstate-cache for fast cached builds"""
+        sh.mkdir("-p", f"build/{self.firmware_machine}/sstate-cache")
+
+    def create_sources_dir(self):
+        """Creates the directory sources for mirror urls used for do_fetch during bitbake build"""
+        sh.mkdir("-p", "build/sources")
+
+    def fetch_sstate_cache(self):
+        """Fetches the sstate_cache from Google Cloud Storage"""
+        cmd = "rsync"
+        bucket = f"gs://{self.gcs_bucket}/cache-unzip/{self.firmware_machine}"
+        sstate_cache_dir = f"build/{self.firmware_machine}/sstate-cache"
+
+        self.gsutil(cmd, bucket, sstate_cache_dir)
+
+    def fetch_sources(self):
+        """Fetches the mirror url sources from Google Cloud Storage"""
+        cmd = "rsync"
+        bucket = f"gs://{self.gcs_bucket}/sources"
+        sources_dir = "build/sources"
+
+        self.gsutil(cmd, bucket, sources_dir)
+
+    def chown_sstate_cache(self):
+        """Changes ownership of sstate-cache directory to user build"""
+        sh.chown("-R", "build", f"build/{self.firmware_machine}/sstate-cache")
+
+    def run_script(self):
+        if self.use_cache:
+            self.create_sstate_cache_dir()
+            self.fetch_sstate_cache()
+            self.chown_sstate_cache()
+
+        self.create_sources_dir()
+        self.fetch_sources()
diff --git a/steps/machine_setup.py b/steps/machine_setup.py
index 49d6062..504600a 100644
--- a/steps/machine_setup.py
+++ b/steps/machine_setup.py
@@ -1,6 +1,6 @@
 from steps import step
-import sh
 import subprocess
+import sh
 
 class MachineSetup(step.Step):
     def run_script(self):
diff --git a/steps/step.py b/steps/step.py
index 820d8c6..80bc8a6 100644
--- a/steps/step.py
+++ b/steps/step.py
@@ -5,14 +5,31 @@
     def __init__(self):
         self.firmware_machine = os.getenv("FIRMWARE_MACHINE")
         self.main_repo_name = os.getenv("MAIN_REPO_NAME")
-        self.repos = os.getenv('REPOS')
-        self.build_script_dir = os.getenv('BUILD_SCRIPT_DIR')
-
-        sh.cd(self.build_script_dir)
+        self.repos = os.getenv("REPOS")
+        self.build_script_dir = os.getenv("BUILD_SCRIPT_DIR")
+        self.use_cache = os.getenv("USE_CACHE") == "True"
+        self.gcs_bucket = os.getenv("GCS_BUCKET")
+        self.firmware_version = os.getenv("FIRMWARE_VERSION")
 
     def parse_repos(self, repos_string):
-        '''Parses a string representation of an array of triples'''
+        """Parses a string representation of an array of triples"""
         return eval(repos_string)
 
+    def export_machine(self):
+        """Exports firmware machine to be used by the build"""
+        return f"export MACHINE={self.firmware_machine}"
+
+    def source_setup(self):
+        """Sets up the firmware machine"""
+        return f"source setup {self.firmware_machine}"
+
+    def umask(self):
+        """Apply umask permissions"""
+        return "umask 002"
+
+    def gsutil(self, cmd, src, dst):
+        """Executes the gsutil tool with cmd from src to dst"""
+        sh.sh("-c", f"gsutil -mq {cmd} -r {src} {dst}")
+
     def run_script(self):
         pass
diff --git a/unit_tests/test_gsutil_setup.py b/unit_tests/test_gsutil_setup.py
new file mode 100644
index 0000000..4704b48
--- /dev/null
+++ b/unit_tests/test_gsutil_setup.py
@@ -0,0 +1,113 @@
+import sh
+import unittest
+from unittest.mock import MagicMock
+from steps import gsutil_setup
+
+
+class TestGsutilSetup(unittest.TestCase):
+    def setUp(self):
+        self.init_step_mock()
+
+        self.step.use_cache = True
+        self.step.create_sstate_cache_dir = MagicMock()
+        self.step.create_sources_dir = MagicMock()
+        self.step.fetch_sstate_cache = MagicMock()
+        self.step.fetch_sources = MagicMock()
+        self.step.chown_sstate_cache = MagicMock()
+
+        self.step.run_script()
+
+    def init_step_mock(self):
+        self.step = gsutil_setup.GsutilSetup()
+        self.step.firmware_machine = "machine0"
+        self.step.gcs_bucket = "a-bucket"
+
+    def test_create_sstate_cache_dir(self):
+        """Tests that sstate_cache_dir method is called"""
+        self.step.create_sstate_cache_dir.assert_called_with()
+
+    def test_create_sources_dir(self):
+        """Tests that create_sources method is called"""
+        self.step.create_sources_dir.assert_called_with()
+
+    def test_fetch_sstate_cache(self):
+        """Tests that fetch_sstate_cache method is called"""
+        self.step.fetch_sstate_cache.assert_called_with()
+
+    def test_fetch_sources(self):
+        """Tests that fetch_sources method is called"""
+        self.step.fetch_sources.assert_called_with()
+
+    def test_disable_sstate_cache(self):
+        """Tests the program state when use_cache=False"""
+        self.init_step_mock()
+
+        self.step.use_cache = False
+        self.step.create_sstate_cache_dir = MagicMock()
+        self.step.create_sources_dir = MagicMock()
+        self.step.fetch_sstate_cache = MagicMock()
+        self.step.fetch_sources = MagicMock()
+        self.step.chown_sstate_cache = MagicMock()
+
+        self.step.run_script()
+
+        self.step.create_sstate_cache_dir.assert_not_called()
+        self.step.create_sources_dir.assert_called_with()
+        self.step.fetch_sstate_cache.assert_not_called()
+        self.step.fetch_sources.assert_called_with()
+        self.step.chown_sstate_cache.assert_not_called()
+
+    def test_mkdir(self):
+        """Tests the call order of the mkdir calls"""
+        self.init_step_mock()
+        self.step.fetch_sstate_cache = MagicMock()
+        self.step.fetch_sources = MagicMock()
+        self.step.chown_sstate_cache = MagicMock()
+        self.step.use_cache = True
+        sh.mkdir = MagicMock()
+
+        self.step.run_script()
+
+        actual_calls = sh.mkdir.call_args_list
+        expected_calls = [(("-p", "build/machine0/sstate-cache"),),
+                    (("-p", "build/sources"),),]
+
+        self.assertEqual(actual_calls, expected_calls)
+
+    def test_gsutil(self):
+        """Tests the order in which gsutil is called, and their arguments"""
+        self.init_step_mock()
+        self.step.use_cache = True
+        self.step.create_sstate_cache_dir = MagicMock()
+        self.step.create_sources_dir = MagicMock()
+        self.step.gsutil = MagicMock()
+        self.step.chown_sstate_cache = MagicMock()
+
+        self.step.run_script()
+
+        actual_calls = self.step.gsutil.call_args_list
+        expected_calls = [(("rsync", "gs://a-bucket/cache-unzip/machine0",
+                    "build/machine0/sstate-cache"),),
+                    (("rsync", "gs://a-bucket/sources", "build/sources"),)]
+
+        self.assertEqual(actual_calls, expected_calls)
+
+    def test_chown(self):
+        """Tests that sh.chown is called with the right arguments when use_cache=True"""
+        self.init_step_mock()
+        self.step.use_cache = True
+        self.step.create_sstate_cache_dir = MagicMock()
+        self.step.create_sources_dir = MagicMock()
+        self.step.fetch_sstate_cache = MagicMock()
+        self.step.fetch_sources = MagicMock()
+        sh.chown = MagicMock()
+
+        self.step.run_script()
+
+        sh.chown("-R", "build", "build/machine0/sstate-cache")
+
+        self.step.chown_sstate_cache = MagicMock()
+
+        self.step.run_script()
+
+        self.step.chown_sstate_cache.assert_called_with()
diff --git a/unit_tests/test_machine_setup.py b/unit_tests/test_machine_setup.py
new file mode 100644
index 0000000..77efd03
--- /dev/null
+++ b/unit_tests/test_machine_setup.py
@@ -0,0 +1,51 @@
+import subprocess
+import sh
+import unittest
+from unittest.mock import MagicMock
+from steps import machine_setup
+
+class TestMachineSetup(unittest.TestCase):
+    def setUp(self):
+        self.init_step_mock()
+        self.step.export_machine = MagicMock(return_value="")
+        self.step.source_setup = MagicMock(return_value="")
+        self.step.umask = MagicMock(return_value="")
+
+        self.step.run_script()
+
+    def init_step_mock(self):
+        self.step = machine_setup.MachineSetup()
+        self.step.firmware_machine = "machine0"
+        self.step.main_repo_name = "repo0"
+        sh.cd = MagicMock()
+        subprocess.run = MagicMock()
+
+    def test_export_machine(self):
+        """Tests that export_machine method is called"""
+        self.step.export_machine.assert_called_with()
+
+    def test_source_setup(self):
+        """Tests that source_setup method is called"""
+        self.step.source_setup.assert_called_with()
+
+    def test_umask(self):
+        """Tests that umask method is called"""
+        self.step.umask.assert_called_with()
+
+    def test_cd(self):
+        """Tests that run_script changes into the main_repo_name dir"""
+        sh.cd.assert_called_with(self.step.main_repo_name)
+
+    def test_subprocess(self):
+        """Tests the arguments of subprocess, and the order in which each
+        sub-commands are called (separated by &&)
+        """
+        self.init_step_mock()
+
+        self.step.run_script()
+
+        command = "export MACHINE=machine0 " \
+                  "&& source setup machine0 " \
+                  "&& umask 002"
+
+        subprocess.run.assert_called_with(["bash", "-c", command])