add common for tests

software/end-to-end-testing/tests/amari_test_common.py

+import json
+import time
+import slapos.testing.e2e as e2e
+from websocket import create_connection
+
+class WebsocketTestBase(e2e.EndToEndTestCase):
+    @classmethod
+    def setUpClass(cls):
+        try:
+            super(WebsocketTestBase, cls).setUpClass()
+            cls.setTestParameters()
+
+            cls.max_retries = 10
+            cls.retry_delay = 180  # seconds
+
+            cls.setup_instances()
+
+            cls.waitUntilGreen(cls.enb_instance_name)
+            cls.waitUntilGreen(cls.cn_instance_name)
+        except Exception as e:
+            cls.logger.error(f"Error during setup: {e}")
+            cls.tearDownClass()
+            raise
+
+    @classmethod
+    def setTestParameters(cls):
+        # This method should be overridden in the subclass to set specific test parameters
+        raise NotImplementedError("You must implement setTestParameters method in the subclass")
+
+    @classmethod
+    def retry_request(cls, func, *args, **kwargs):
+        for attempt in range(cls.max_retries):
+            try:
+                result = func(*args, **kwargs)
+                if result:
+                    return result
+            except Exception as e:
+                cls.logger.error(f"Error on attempt {attempt + 1}: {e}")
+            if attempt < cls.max_retries - 1:
+                time.sleep(cls.retry_delay)
+        return None
+
+    @classmethod
+    def setup_instances(cls):
+        cls.request_enb()
+        cls.request_core_network()
+        cls.setup_websocket_connection()
+
+    @classmethod
+    def request_enb(cls, custom_params=None):
+        cls.logger.info(f"Requesting eNB: {cls.enb_instance_name}")
+        enb_parameters = {
+            "dl_earfcn": cls.dl_earfcn,
+            "plmn_list": {"Australia": {"plmn": "50501"}}
+        }
+
+        if custom_params:
+            enb_parameters.update(custom_params)
+
+        json_enb_parameters = json.dumps(enb_parameters)
+        cls.retry_request(cls.request, cls.product, cls.enb_instance_name,
+                          filter_kw={"computer_guid": cls.comp_enb},
+                          partition_parameter_kw={'_': json_enb_parameters},
+                          software_type='enb')
+
+    @classmethod
+    def request_core_network(cls):
+        cls.logger.info(f"Requesting Core Network: {cls.cn_instance_name}")
+        core_network_parameters = json.dumps({"core_network_plmn": "50501"})
+        cls.retry_request(cls.request_core_network_with_guid, core_network_parameters)
+
+    @classmethod
+    def request_core_network_with_guid(cls, core_network_parameters):
+        core_network_instance = cls.request(cls.product, cls.cn_instance_name,
+                                            filter_kw={"computer_guid": cls.comp_cn},
+                                            partition_parameter_kw={'_': core_network_parameters},
+                                            software_type='core-network')
+        if core_network_instance:
+            instance_infos = cls.getInstanceInfos(cls.cn_instance_name)
+            cls.cn_instance_guid = instance_infos.news['instance'][0]['reference']
+            cls.request_demo_sim_cards()
+            return True
+        return False
+
+    @classmethod
+    def request_demo_sim_cards(cls):
+        if cls.cn_instance_guid is None:
+            cls.logger.error("Core network instance GUID not set. Cannot request demo SIM cards.")
+            return
+
+        cls.logger.info(f"Requesting Demo SIM Cards: {cls.sim_instance_name}")
+        sim_card_parameters = {
+            "sim_algo": "xor",
+            "imsi": "505010123456789",
+            "k": "00112233445566778899aabbccddeeff",
+            "imeisv": "8682430000000101",
+            "impi": "505010123456789@ims.mnc505.mcc001.3gppnetwork.org",
+            "impu": ["505010123456789", "tel:0600000000", "tel:600"]
+        }
+
+        json_sim_card_parameters = json.dumps(sim_card_parameters)
+        cls.retry_request(cls.request, cls.product, cls.sim_instance_name,
+                          partition_parameter_kw={'_': json_sim_card_parameters},
+                          software_type='core-network',
+                          filter_kw={"instance_guid": cls.cn_instance_guid},
+                          shared=True, state='started')
+
+    @classmethod
+    def setup_websocket_connection(cls):
+        ue_instance = cls.retry_request(cls.request_ue)
+        cls.waitUntilGreen(cls.ue_instance_name)
+        cls.ue_com_addr = ue_instance.get('com_addr') if ue_instance else None
+        if not cls.ue_com_addr:
+            cls.logger.error("Failed to obtain UE com address.")
+            return
+
+        cls.ws_url = f"ws://{cls.ue_com_addr}"
+        cls.logger.info(f"Websocket URL: {cls.ws_url}")
+
+        for attempt in range(cls.max_retries):
+            try:
+                cls.ws = create_connection(cls.ws_url)
+                cls.logger.info("Websocket connection established.")
+                break
+            except Exception as e:
+                cls.logger.error(f"Websocket connection attempt {attempt + 1} failed: {e}")
+                if attempt < cls.max_retries - 1:
+                    time.sleep(5)
+
+    @classmethod
+    def request_ue(cls):
+        ue_parameters = {
+            "n_antenna_dl": 2,
+            "n_antenna_ul": 2,
+            "dl_earfcn": cls.dl_earfcn,
+            "sim_algo": "xor",
+            "imsi": "505010123456789",
+            "k": "00112233445566778899aabbccddeeff",
+            "imeisv": "8682430000000101",
+            "impi": "505010123456789@ims.mnc505.mcc001.3gppnetwork.org",
+            "impu": ["505010123456789", "tel:0600000000", "tel:600"]
+        }
+        json_ue_parameters = json.dumps(ue_parameters)
+        return cls.retry_request(cls.request, cls.ue_product, cls.ue_instance_name,
+                                 filter_kw={"computer_guid": cls.comp_ue},
+                                 partition_parameter_kw={'_': json_ue_parameters},
+                                 software_type='ue-lte')
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, 'ws') and cls.ws:
+            cls.ws.close()
+        super(WebsocketTestBase, cls).tearDownClass()
+
+    def send(self, msg):
+        self.ws.send(json.dumps(msg))
+
+    def recv(self):
+        return json.loads(self.ws.recv())
+
+
+    def ue_get(self):
+        self.send({"message": "ue_get"})
+        result = self.recv()
+
+        if 'message' not in result:
+            raise ValueError(f"Unexpected response format: {result}")
+
+        if 'ue_list' in result:
+            if not result['ue_list']:
+                raise ValueError(f"No UE found in response: {result}")
+            return result['ue_list'][0]
+        else:
+            return result
+
+
+    def power_on(self, ue_id):
+        self.assertFalse(self.ue_get()['power_on'], "UE already powered on")
+        self.send({"message": "power_on", "ue_id": ue_id})
+        self.recv()
+
+    def power_off(self, ue_id):
+        self.assertTrue(self.ue_get()['power_on'], "UE already powered off")
+        self.send({"message": "power_off", "ue_id": ue_id})
+        self.recv()