From eeb5c187b09e2b9e6fd476b9012c3e31fb5f6424 Mon Sep 17 00:00:00 2001
From: Dheyay <dheyay.desai@canonical.com>
Date: Wed, 6 Mar 2024 13:07:58 -0800
Subject: [PATCH] Add coverage for behave tests

Fixes: #2903
---
 features/environment.py    | 53 ++++++++++++++++++++++++++++++++++++++
 features/steps/machines.py | 19 +++++++++++++-
 features/steps/shell.py    | 13 +++++++++-
 3 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/features/environment.py b/features/environment.py
index 2d7dd69449..4170d534cf 100644
--- a/features/environment.py
+++ b/features/environment.py
@@ -412,6 +412,59 @@ def before_scenario(context: Context, scenario: Scenario):
             )
 
 
+def after_scenario(context, scenario):
+    """Collect the coverage files after the scenario is run."""
+    cov_dir = os.path.join(context.pro_config.artifact_dir, "coverage")
+    if not os.path.exists(cov_dir):
+        os.makedirs(cov_dir)
+
+    inner_dir = os.path.join(
+        datetime.datetime.now().strftime("%Y-%m-%d"),
+        "{}".format(os.path.basename(scenario.filename.replace(".", "_"))),
+    )
+    new_artifacts_dir = os.path.join(
+        cov_dir,
+        inner_dir,
+    )
+    if not os.path.exists(new_artifacts_dir):
+        os.makedirs(new_artifacts_dir)
+    if hasattr(context, "machines") and SUT in context.machines:
+        try:
+            scenario_name = (
+                os.path.basename(scenario.filename.replace(".", "_"))
+                + "_"
+                + str(scenario.line)
+            )
+            cov_filename = ".coverage.{}".format(scenario_name)
+            context.machines[SUT].instance.execute(
+                [
+                    "bash",
+                    "-c",
+                    "mv .coverage /tmp/{cov_filename}".format(
+                        cov_filename=cov_filename
+                    ),
+                ],
+                use_sudo=True,
+            )
+            context.machines[SUT].instance.execute(
+                [
+                    "chmod",
+                    "666",
+                    "/tmp/{cov_filename}".format(cov_filename=cov_filename),
+                ],
+                use_sudo=True,
+            )
+
+            dest = os.path.join(new_artifacts_dir, cov_filename)
+            context.machines[SUT].instance.pull_file(
+                "/tmp/{cov_filename}".format(cov_filename=cov_filename), dest
+            )
+            logging.warning("Done collecting coverage.")
+        except Exception as e:
+            logging.error(str(e))
+            logging.warning("Failed to collect coverage")
+
+
 def after_step(context, step):
     """Collect test artifacts in the event of failure."""
     if step.status == "failed":
diff --git a/features/steps/machines.py b/features/steps/machines.py
index 80d4e39307..44e5fee318 100644
--- a/features/steps/machines.py
+++ b/features/steps/machines.py
@@ -6,7 +6,7 @@
 from behave import given, when
 from pycloudlib.instance import BaseInstance  # type: ignore
 
-from features.steps.packages import when_i_apt_update
+from features.steps.packages import when_i_apt_install, when_i_apt_update
 from features.steps.shell import when_i_run_command
 from features.steps.ubuntu_advantage_tools import when_i_install_uat
 from features.util import (
@@ -110,6 +110,23 @@ def given_a_machine(
     # make sure the machine has up-to-date apt data
     when_i_apt_update(context, machine_name=machine_name)
 
+    # add coverage
+    when_i_apt_install(context, "python3-coverage", machine_name=machine_name)
+
+    # Create the .coveragerc file in the lxd container
+    if hasattr(context, "machines") and SUT in context.machines:
+        context.machines[SUT].instance.execute(
+            [
+                "bash",
+                "-c",
+                "echo -e '[run]\nrelative_files = True\n\n[paths]\nsource = \
+                    \n\tuaclient/ \
+                    \n\tusr/lib/python3/dist-packages/uaclient/' > \
+                    /usr/lib/python3/dist-packages/uaclient/.coveragerc",
+            ],
+            use_sudo=True,
+        )
+
     if cleanup:
 
         def cleanup_instance():
diff --git a/features/steps/shell.py b/features/steps/shell.py
index a8b67120e9..d6045a2ca4 100644
--- a/features/steps/shell.py
+++ b/features/steps/shell.py
@@ -64,7 +64,18 @@ def when_i_run_command(
 ):
     command = process_template_vars(context, command)
     prefix = get_command_prefix_for_user_spec(user_spec)
-    full_cmd = prefix + shlex.split(command)
+    if "pro" in command.split():
+        command = command.replace(
+            "pro",
+            "python3 -m coverage run \
+                --rcfile=/usr/lib/python3/dist-packages/uaclient/.coveragerc \
+                --source=/usr/lib/python3/dist-packages/uaclient /usr/bin/pro",
+            1,
+        )
+
+    split_command = shlex.split(command)
+    print(split_command)
+    full_cmd = prefix + split_command
 
     if stdin is not None:
         stdin = stdin.replace("\\n", "\n")