Chef - Add BUILD.gn and unit tests for stateful_shell.py (#19205)

Author: cpagravel

BUILD.gn

@@ -166,6 +166,7 @@ if (current_toolchain != "${dir_pw_toolchain}/default:default") {
     if (chip_link_tests) {
       deps = [
         "//:fake_platform_tests",
+        "//examples/chef:chef.tests",
         "//scripts/build:build_examples.tests",
         "//scripts/idl:idl.tests",
         "//src:tests_run",

examples/chef/BUILD.gn

@@ -0,0 +1,32 @@
+# Copyright (c) 2022 Project CHIP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import("//build_overrides/build.gni")
+import("//build_overrides/chip.gni")
+
+import("//build_overrides/pigweed.gni")
+import("$dir_pw_build/python.gni")
+
+pw_python_package("chef") {
+  setup = [ "setup.py" ]
+
+  sources = [
+    "__init__.py",
+    "chef.py",
+    "constants.py",
+    "stateful_shell.py",
+  ]
+
+  tests = [ "test_stateful_shell.py" ]
+}

examples/chef/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2022 Project CHIP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""The chef package."""
+
+import setuptools  # type: ignore
+
+setuptools.setup(
+    name='chef',
+    version='0.0.1',
+    author='Project CHIP Authors',
+    description='Build custom sample apps for supported platforms',
+    packages=setuptools.find_packages(),
+    package_data={'chef': ['py.typed']},
+    zip_safe=False,
+)

examples/chef/setup.py

@@ -16,19 +16,26 @@
 import shlex
 import subprocess
 import sys
+import time
 from typing import Dict, Optional
 
 import constants
 
 _ENV_FILENAME = ".shell_env"
 _OUTPUT_FILENAME = ".shell_output"
 _HERE = os.path.dirname(os.path.abspath(__file__))
+_TEE_WAIT_TIMEOUT = 3
 
 TermColors = constants.TermColors
 
 
 class StatefulShell:
-    """A Shell that tracks state changes of the environment."""
+    """A Shell that tracks state changes of the environment.
+
+    Attributes:
+        env: Env variables passed to command. It gets updated after every command.
+        cwd: Current working directory of shell.
+    """
 
     def __init__(self) -> None:
         if sys.platform == "linux" or sys.platform == "linux2":
@@ -44,8 +51,8 @@ def __init__(self) -> None:
 
         # This file holds the env after running a command. This is a better approach
         # than writing to stdout because commands could redirect the stdout.
-        self.envfile_path: str = os.path.join(_HERE, _ENV_FILENAME)
-        self.cmd_output_path: str = os.path.join(_HERE, _OUTPUT_FILENAME)
+        self._envfile_path: str = os.path.join(_HERE, _ENV_FILENAME)
+        self._cmd_output_path: str = os.path.join(_HERE, _OUTPUT_FILENAME)
 
     def print_env(self) -> None:
         """Print environment variables in commandline friendly format for export.
@@ -87,13 +94,25 @@ def run_cmd(
         if return_cmd_output:
             # Piping won't work here because piping will affect how environment variables
             # are propagated. This solution uses tee without piping to preserve env variables.
-            redirect = f" > >(tee \"{self.cmd_output_path}\") 2>&1 "  # include stderr
+            redirect = f" > >(tee \"{self._cmd_output_path}\") 2>&1 "  # include stderr
+
+            # Delete the file before running the command so we can later check if the file
+            # exists as a signal that tee has finished writing to the file.
+            if os.path.isfile(self._cmd_output_path):
+                os.remove(self._cmd_output_path)
         else:
             redirect = ""
 
+        # TODO: Use env -0 when `macos-latest` refers to macos-12 in github actions.
+        # env -0 is ideal because it will support cases where an env variable that has newline
+        # characters. The flag "-0" is requires MacOS 12 which is still in beta in Github Actions.
+        # The less ideal `env` command is used by itself, with the caveat that newline chars
+        # are unsupported in env variables.
+        save_env_cmd = f"env > {self._envfile_path}"
+
         command_with_state = (
-            f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd} {redirect}; RETCODE=$?;"
-            f" env -0 > {self.envfile_path}; exit $RETCODE")
+            f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd} {redirect}; RETCODE=$?; "
+            f"{save_env_cmd}; exit $RETCODE")
         with subprocess.Popen(
             [command_with_state],
             env=self.env, cwd=self.cwd,
@@ -102,9 +121,9 @@ def run_cmd(
             returncode = proc.wait()
 
         # Load env state from envfile.
-        with open(self.envfile_path, encoding="latin1") as f:
-            # Split on null char because we use env -0.
-            env_entries = f.read().split("\0")
+        with open(self._envfile_path, encoding="latin1") as f:
+            # TODO: Split on null char after updating to env -0 - requires MacOS 12.
+            env_entries = f.read().split("\n")
             for entry in env_entries:
                 parts = entry.split("=")
                 # Handle case where an env variable contains text with '='.
@@ -119,6 +138,21 @@ def run_cmd(
                 f"\nCmd: {cmd}")
 
         if return_cmd_output:
-            with open(self.cmd_output_path, encoding="latin1") as f:
-                output = f.read()
+            # Poll for file due to give 'tee' time to close.
+            # This is necessary because 'tee' waits for all subshells to finish before writing.
+            start_time = time.time()
+            while time.time() - start_time < _TEE_WAIT_TIMEOUT:
+                try:
+                    with open(self._cmd_output_path, encoding="latin1") as f:
+                        output = f.read()
+                    break
+                except FileNotFoundError:
+                    pass
+                time.sleep(0.1)
+            else:
+                raise TimeoutError(
+                    f"Error. Output file: {self._cmd_output_path} not created within "
+                    f"the alloted time of: {_TEE_WAIT_TIMEOUT}s"
+                )
+
             return output

examples/chef/stateful_shell.py

@@ -0,0 +1,48 @@
+"""Tests for stateful_shell.py
+
+Usage:
+python -m unittest
+"""
+
+import unittest
+
+import stateful_shell
+
+
+class TestStatefulShell(unittest.TestCase):
+    """Testcases for stateful_shell.py."""
+
+    def setUp(self):
+        """Prepares stateful shell instance for tests."""
+        self.shell = stateful_shell.StatefulShell()
+
+    def test_cmd_output(self):
+        """Tests shell command output."""
+        resp = self.shell.run_cmd("echo test123", return_cmd_output=True).strip()
+        self.assertEqual(resp, "test123")
+
+    def test_set_env_in_shell(self):
+        """Tests setting env variables in shell."""
+        self.shell.run_cmd("export TESTVAR=123")
+        self.assertEqual(self.shell.env["TESTVAR"], "123")
+
+    def test_set_env_outside_shell(self):
+        """Tests setting env variables outside shell call."""
+        self.shell.env["TESTVAR"] = "1234"
+        resp = self.shell.run_cmd("echo $TESTVAR", return_cmd_output=True).strip()
+        self.assertEqual(resp, "1234")
+
+    def test_env_var_set_get(self):
+        """Tests setting and getting env vars between calls."""
+        self.shell.run_cmd("export TESTVAR=123")
+        resp = self.shell.run_cmd("echo $TESTVAR", return_cmd_output=True).strip()
+        self.assertEqual(resp, "123")
+
+    def test_raise_on_returncode(self):
+        """Tests raising errors when returncode is nonzero."""
+        with self.assertRaises(RuntimeError):
+            self.shell.run_cmd("invalid_cmd > /dev/null 2>&1", raise_on_returncode=True)
+
+
+if __name__ == "__main__":
+    unittest.main()