From bdd56dc60be2fa26f59f7a9d46b64505872710e3 Mon Sep 17 00:00:00 2001
From: Vinnie Magro <vmagro@meta.com>
Date: Thu, 3 Oct 2024 10:02:12 -0700
Subject: [PATCH] [antlir2][rpm] upgrade should install if missing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Summary: If an rpm is not installed, upgrade should install it, not fail

Test Plan:
```
❯ buck2 test fbcode//antlir/antlir2/features/rpm/tests:upgrade-not-already-installed
Buck UI: https://www.internalfb.com/buck2/d0351215-69da-4675-8529-3bb752bf4574
Test UI: https://www.internalfb.com/intern/testinfra/testrun/7036874653273252
Tests finished: Pass 1. Fail 0. Fatal 0. Skip 0. Build failure 0
```

Reviewed By: naveedgol

Differential Revision: D63806983

fbshipit-source-id: ebcba322cfa568de61a7fba4fc8dfc59f8b85a4c
---
 antlir/antlir2/features/rpm/driver.py  |  9 ++++++++-
 antlir/antlir2/features/rpm/tests/BUCK | 26 ++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/antlir/antlir2/features/rpm/driver.py b/antlir/antlir2/features/rpm/driver.py
index cc91f0ac3b..630fb640b9 100755
--- a/antlir/antlir2/features/rpm/driver.py
+++ b/antlir/antlir2/features/rpm/driver.py
@@ -248,13 +248,20 @@ def resolve(out, spec, base, local_rpms, explicitly_installed_package_names):
                         json.dump({"package_not_found": e.pkg_spec}, o)
         elif action == "upgrade":
             if isinstance(source, dnf.package.Package):
-                base.package_upgrade(source)
+                try:
+                    base.package_upgrade(source)
+                except dnf.exceptions.MarkingError:
+                    # If it's not installed, upgrade should behave the same as install
+                    base.package_install(source, strict=True)
             else:
                 try:
                     base.upgrade(source)
                 except dnf.exceptions.PackageNotFoundError as e:
                     with out as o:
                         json.dump({"package_not_found": e.pkg_spec}, o)
+                except dnf.exceptions.PackagesNotInstalledError:
+                    # If it's not installed, upgrade should behave the same as install
+                    base.install(source, strict=True)
         elif action == "remove":
             # cannot remove by file path, so let's do this to be extra safe
             try:
diff --git a/antlir/antlir2/features/rpm/tests/BUCK b/antlir/antlir2/features/rpm/tests/BUCK
index e7408284b9..506884121f 100644
--- a/antlir/antlir2/features/rpm/tests/BUCK
+++ b/antlir/antlir2/features/rpm/tests/BUCK
@@ -171,6 +171,32 @@ test_rpms(
     parent_layer = simple,
 )
 
+test_rpms(
+    name = "upgrade-not-already-installed",
+    expected = expected_t(
+        installed = [
+            "foo-3-1",
+        ],
+    ),
+    features = [
+        feature.rpms_upgrade(rpms = ["foo"]),
+        ":test-deps",
+    ],
+)
+
+test_rpms(
+    name = "upgrade-not-already-installed-file",
+    expected = expected_t(
+        installed = [
+            "foo-2-1",
+        ],
+    ),
+    features = [
+        feature.rpms_upgrade(rpms = ["//antlir/antlir2/features/rpm/tests/repo:foo-2-1.noarch"]),
+        ":test-deps",
+    ],
+)
+
 test_rpms(
     name = "downgrade",
     expected = expected_t(