From d0c7174666e4ef9807165974bcedfc91ff8c73fd Mon Sep 17 00:00:00 2001
From: Sean Marlow <sean.marlow@suse.com>
Date: Thu, 30 Nov 2023 16:13:18 -0500
Subject: [PATCH] Handle suma payg in ec2 billing code test

---
 .../SLES/EC2/test_sles_ec2_billing_code.py    | 11 +++++-
 usr/share/lib/img_proof/tests/conftest.py     | 37 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/usr/share/lib/img_proof/tests/SLES/EC2/test_sles_ec2_billing_code.py b/usr/share/lib/img_proof/tests/SLES/EC2/test_sles_ec2_billing_code.py
index 0f5a25b..90cb58b 100644
--- a/usr/share/lib/img_proof/tests/SLES/EC2/test_sles_ec2_billing_code.py
+++ b/usr/share/lib/img_proof/tests/SLES/EC2/test_sles_ec2_billing_code.py
@@ -4,7 +4,8 @@
 def test_sles_ec2_billing_code(
     get_ec2_billing_products,
     is_byos,
-    get_release_value
+    get_release_value,
+    is_suma_server
 ):
     try:
         products = get_ec2_billing_products()
@@ -17,7 +18,13 @@ def test_sles_ec2_billing_code(
     variant = get_release_value('VARIANT_ID')
     byos = is_byos()
 
-    if byos or (variant in ('sles-sap', 'sles-sap-hardened') and not byos):
+    has_no_code = (
+        byos or
+        (variant in ('sles-sap', 'sles-sap-hardened') and not byos) or
+        is_suma_server() and not byos
+    )
+
+    if has_no_code:
         assert products is None
     else:
         assert products
diff --git a/usr/share/lib/img_proof/tests/conftest.py b/usr/share/lib/img_proof/tests/conftest.py
index d81482c..0907ff1 100644
--- a/usr/share/lib/img_proof/tests/conftest.py
+++ b/usr/share/lib/img_proof/tests/conftest.py
@@ -145,6 +145,43 @@ def f():
     return f
 
 
+@pytest.fixture()
+def is_suma_server(host, get_baseproduct):
+    def f():
+        suma_server_product = '/etc/products.d/SUSE-Manager-Server.prod'
+        suma_server = host.file(suma_server_product)
+        base_product = get_baseproduct()
+
+        return all([
+            suma_server.exists,
+            suma_server.is_file,
+            base_product == suma_server_product
+        ])
+    return f
+
+
+@pytest.fixture()
+def is_suma_proxy(host, get_baseproduct):
+    def f():
+        suma_proxy_product = '/etc/products.d/SUSE-Manager-Proxy.prod'
+        suma_proxy = host.file(suma_proxy_product)
+        base_product = get_baseproduct()
+
+        return all([
+            suma_proxy.exists,
+            suma_proxy.is_file,
+            base_product == suma_proxy_product
+        ])
+    return f
+
+
+@pytest.fixture()
+def is_suma(is_suma_server, is_suma_proxy):
+    def f():
+        return is_suma_server() or is_suma_proxy()
+    return f
+
+
 @pytest.fixture()
 def is_sles_sapcal(host):
     def f():