register function made more flexible

(cherry picked from commit ceb0241) Signed-off-by: Gaurav Talreja <[email protected]>
SatelliteQE · Aug 7, 2023 · 4c93d25 · 4c93d25
1 parent c1b5c7d
commit 4c93d25
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 7 deletions.
diff --git a/robottelo/hosts.py b/robottelo/hosts.py
@@ -563,8 +563,8 @@ def register(
         using a global registration template.
 
         :param target: Satellite or Capusle object to register to, required.
-        :param org: Organization to register content host for, required.
-        :param loc: Location to register content host for, required.
+        :param org: Organization to register content host to. Previously required, pass None to omit
+        :param loc: Location to register content host for, Previously required, pass None to omit.
         :param activation_keys: Activation key name to register content host with, required.
         :param setup_insights: Install and register Insights client, requires OS repo.
         :param setup_remote_execution: Copy remote execution SSH key.
@@ -584,11 +584,25 @@ def register(
         """
         options = {
             'activation-keys': activation_keys,
-            'organization-id': org.id,
-            'location-id': loc.id,
             'insecure': str(insecure).lower(),
             'update-packages': str(update_packages).lower(),
         }
+        if org is not None:
+            if isinstance(org, entities.Organization):
+                options['organization-id'] = org.id
+            elif isinstance(org, dict):
+                options['organization-id'] = org['id']
+            else:
+                raise ValueError('org must be a dict or an Organization object')
+
+        if loc is not None:
+            if isinstance(loc, entities.Location):
+                options['location-id'] = loc.id
+            elif isinstance(loc, dict):
+                options['location-id'] = loc['id']
+            else:
+                raise ValueError('loc must be a dict or a Location object')
+
         if target.__class__.__name__ == 'Capsule':
             options['smart-proxy'] = target.hostname
         elif target is not None and target.__class__.__name__ not in ['Capsule', 'Satellite']:

diff --git a/tests/foreman/longrun/test_oscap.py b/tests/foreman/longrun/test_oscap.py
@@ -198,12 +198,13 @@ def test_positive_oscap_run_via_ansible(
         }
     )
     with Broker(nick=distro, host_class=ContentHost, deploy_flavor=settings.flavors.default) as vm:
-        result = vm.register(module_org, None, ak_name[distro], target_sat)
-        assert result.status == 0, f'Failed to register host: {result.stderr}'
         if distro not in ('rhel7'):
             vm.create_custom_repos(**rhel_repo)
         else:
             vm.create_custom_repos(**{distro: rhel_repo})
+        result = vm.register(module_org, None, ak_name[distro], target_sat)
+        assert result.status == 0, f'Failed to register host: {result.stderr}'
+
         Host.update(
             {
                 'name': vm.hostname.lower(),
@@ -297,9 +298,9 @@ def test_positive_oscap_run_via_ansible_bz_1814988(
         }
     )
     with Broker(nick='rhel7', host_class=ContentHost, deploy_flavor=settings.flavors.default) as vm:
+        vm.create_custom_repos(rhel7=settings.repos.rhel7_os)
         result = vm.register(module_org, None, ak_name['rhel7'], target_sat)
         assert result.status == 0, f'Failed to register host: {result.stderr}'
-        vm.create_custom_repos(rhel7=settings.repos.rhel7_os)
         # Harden the rhel7 client with DISA STIG security policy
         vm.run('yum install -y scap-security-guide')
         vm.run(