yaml: prefer importing the system yaml module

Check for the existence of yaml in the system libs before going the
Return an error when no module is found.

Using the system yaml guarantees that safe_load and safe_dump works.

LP: #2007234
LP: #2007241

Signed-off-by: Renan Rodrigo <renanrodrigo@canonical.com>

features/unattached_commands.feature

@@ -349,3 +349,39 @@ Feature: Command behaviour when unattached
           | bionic  |
           | focal   |
           | jammy   |
+
+    @series.jammy
+    @series.kinetic
+    @series.lunar
+    @uses.config.machine_type.lxd.container
+    Scenario Outline: services don't fail when an incompatible yaml package is installed
+        Given a `<release>` machine with ubuntu-advantage-tools installed
+        When I run `apt update` with sudo
+        And I run `apt install python3-pip -y` with sudo
+        And I run `pip3 install pyyaml==3.10` with sudo
+        And I run `ls /usr/local/lib/<python_version>/dist-packages/` with sudo
+        Then stdout matches regexp:
+        """
+        yaml
+        """
+        # Test the specific script which triggered LP #2007241
+        When I run `python3 /usr/lib/ubuntu-advantage/esm_cache.py` with sudo
+        And I run `systemctl start esm-cache.service` with sudo
+        And I wait `5` seconds
+        And I run `systemctl --failed` with sudo
+        Then stdout does not match regexp:
+        """
+        esm-cache\.service
+        """        
+        When I run `systemctl start apt-news.service` with sudo
+        And I wait `5` seconds
+        And I run `systemctl --failed` with sudo
+        Then stdout does not match regexp:
+        """
+        apt-news\.service
+        """
+        Examples: ubuntu release
+          | release | python_version |
+          | jammy   | python3.10     |
+          | kinetic | python3.10     |
+          | lunar   | python3.11     |

uaclient/cli.py

@@ -14,8 +14,6 @@
 from functools import wraps
 from typing import List, Optional, Tuple  # noqa
 
-import yaml
-
 from uaclient import (
     actions,
     apt,
@@ -33,7 +31,7 @@
     security_status,
 )
 from uaclient import status as ua_status
-from uaclient import util, version
+from uaclient import util, version, yaml
 from uaclient.api.api import call_api
 from uaclient.api.u.pro.attach.auto.full_auto_attach.v1 import (
     FullAutoAttachOptions,

uaclient/config.py

@@ -6,8 +6,6 @@
 from functools import lru_cache, wraps
 from typing import Any, Callable, Dict, Optional, Tuple, TypeVar
 
-import yaml
-
 from uaclient import (
     apt,
     event_logger,
@@ -17,6 +15,7 @@
     snap,
     system,
     util,
+    yaml,
 )
 from uaclient.defaults import (
     APT_NEWS_URL,

uaclient/entitlements/base.py

@@ -5,9 +5,15 @@
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
-import yaml
-
-from uaclient import config, contract, event_logger, messages, system, util
+from uaclient import (
+    config,
+    contract,
+    event_logger,
+    messages,
+    system,
+    util,
+    yaml,
+)
 from uaclient.defaults import DEFAULT_HELP_FILE
 from uaclient.entitlements.entitlement_status import (
     ApplicabilityStatus,

uaclient/event_logger.py

@@ -11,7 +11,7 @@
 import sys
 from typing import Any, Dict, List, Optional, Set, Union  # noqa: F401
 
-import yaml
+from uaclient import yaml
 
 JSON_SCHEMA_VERSION = "0.1"
 EventFieldErrorType = Optional[Union[str, Dict[str, str]]]

uaclient/files/data_types.py

@@ -2,9 +2,7 @@
 from enum import Enum
 from typing import Callable, Dict, Generic, Optional, Type, TypeVar
 
-import yaml
-
-from uaclient import exceptions
+from uaclient import exceptions, yaml
 from uaclient.data_types import DataObject
 from uaclient.files.files import UAFile
 from uaclient.util import DatetimeAwareJSONDecoder

uaclient/messages.py

@@ -1169,3 +1169,10 @@ class TxtColor:
 
 $ sudo rm {lock_file_path}""",
 )
+
+MISSING_YAML_MODULE = NamedMessage(
+    "missing-yaml-module",
+    """\
+Couldn't import the YAML module from /usr/lib/python3/dist-packages.
+Make sure the 'python3-yaml' package is installed correctly.""",
+)

uaclient/tests/test_cli_attach.py

@@ -6,9 +6,8 @@
 
 import mock
 import pytest
-import yaml
 
-from uaclient import event_logger, messages, status, util
+from uaclient import event_logger, messages, status, util, yaml
 from uaclient.cli import (
     UA_AUTH_TOKEN_URL,
     action_attach,

uaclient/tests/test_cli_status.py

@@ -9,9 +9,8 @@
 
 import mock
 import pytest
-import yaml
 
-from uaclient import exceptions, messages, status, util
+from uaclient import exceptions, messages, status, util, yaml
 from uaclient.cli import action_status, get_parser, main, status_parser
 from uaclient.conftest import FakeNotice
 from uaclient.event_logger import EventLoggerMode

uaclient/tests/test_config.py

@@ -8,9 +8,8 @@
 
 import mock
 import pytest
-import yaml
 
-from uaclient import apt, exceptions, messages, util
+from uaclient import apt, exceptions, messages, util, yaml
 from uaclient.config import (
     PRIVATE_SUBDIR,
     UA_CONFIGURABLE_KEYS,

uaclient/tests/test_event_logger.py

@@ -4,8 +4,8 @@
 
 import mock
 import pytest
-import yaml
 
+from uaclient import yaml
 from uaclient.event_logger import JSON_SCHEMA_VERSION, EventLoggerMode
 
 

uaclient/tests/test_yaml.py

@@ -0,0 +1,60 @@
+import mock
+import pytest
+
+from uaclient.exceptions import UserFacingError
+from uaclient.yaml import get_imported_yaml_module
+
+M_PREFIX = "uaclient.yaml."
+
+
+@mock.patch(M_PREFIX + "sys")
+class TestYamlImport:
+    def test_get_yaml_module_returns_imported(self, m_sys):
+        imported_yaml = mock.MagicMock()
+        m_sys.modules = {"yaml": imported_yaml}
+        yaml = get_imported_yaml_module()
+        assert yaml is imported_yaml
+
+    @pytest.mark.parametrize("has_yaml_in_system", (True, False))
+    @mock.patch(M_PREFIX + "importlib_machinery")
+    @mock.patch(M_PREFIX + "importlib_util")
+    def test_get_yaml_module_imports_from_system(
+        self, m_util, m_machinery, m_sys, has_yaml_in_system
+    ):
+        m_sys.modules = {}
+        m_finder = mock.MagicMock()
+        m_machinery.PathFinder.return_value = m_finder
+
+        m_yaml_spec = mock.MagicMock()
+        m_yaml_module = mock.MagicMock()
+
+        m_util.module_from_spec.return_value = m_yaml_module
+
+        if has_yaml_in_system:
+            m_finder.find_spec.return_value = m_yaml_spec
+        else:
+            m_finder.find_spec.return_value = None
+
+        if has_yaml_in_system:
+            yaml = get_imported_yaml_module()
+
+            assert m_yaml_module == yaml
+            assert [
+                mock.call("yaml", path=["/usr/lib/python3/dist-packages"])
+            ] == m_finder.find_spec.call_args_list
+            assert [
+                mock.call(m_yaml_spec)
+            ] == m_util.module_from_spec.call_args_list
+            assert m_sys.modules["yaml"] == m_yaml_module
+            assert [
+                mock.call(m_yaml_module)
+            ] == m_yaml_spec.loader.exec_module.call_args_list
+        else:
+            with pytest.raises(UserFacingError) as excinfo:
+                yaml = get_imported_yaml_module()
+
+            assert "missing-yaml-module" == excinfo.value.msg_code
+            assert "Couldn't import the YAML module" in excinfo.value.msg
+            assert [
+                mock.call("yaml", path=["/usr/lib/python3/dist-packages"])
+            ] == m_finder.find_spec.call_args_list

uaclient/yaml.py

@@ -0,0 +1,43 @@
+import sys
+from importlib import machinery as importlib_machinery
+from importlib import util as importlib_util
+
+from uaclient.exceptions import UserFacingError
+from uaclient.messages import MISSING_YAML_MODULE
+
+
+def get_imported_yaml_module():
+    # Maybe we already imported it?
+    if "yaml" in sys.modules:
+        return sys.modules["yaml"]
+
+    # Try to get the system yaml module
+    finder = importlib_machinery.PathFinder()
+    pyyaml_spec = finder.find_spec(
+        "yaml", path=["/usr/lib/python3/dist-packages"]
+    )
+    # If there is no system 'yaml' module (which is weird enough),
+    # then raise an exception asking for it to be there
+    if pyyaml_spec is None:
+        raise UserFacingError(
+            MISSING_YAML_MODULE.msg, MISSING_YAML_MODULE.name
+        )
+
+    # Import it!
+    yaml_module = importlib_util.module_from_spec(pyyaml_spec)
+    sys.modules["yaml"] = yaml_module
+    loader = pyyaml_spec.loader
+    # The loader shouldn't be None, documentation says:
+    # "The finder should always set this attribute."
+    # But mypy complains, so we check anyway.
+    if loader is not None:
+        loader.exec_module(yaml_module)
+
+    return yaml_module
+
+
+yaml = get_imported_yaml_module()
+
+safe_load = yaml.safe_load
+safe_dump = yaml.safe_dump
+parser = yaml.parser