diff --git a/src/device/.gitlab-ci.yml b/src/device/.gitlab-ci.yml
index 3ae6b2b20d22d0a91d10b61acd20fd0d2d4a28ac..16fd27444deba125fe25fc07eec701a1278555c4 100644
--- a/src/device/.gitlab-ci.yml
+++ b/src/device/.gitlab-ci.yml
@@ -38,63 +38,65 @@ build device:
       - manifests/${IMAGE_NAME}service.yaml
       - .gitlab-ci.yml
 
-## Start Mock QKD Nodes before unit testing
-#start_mock_nodes:
-#  stage: deploy
-#  script:
-#    - bash src/tests/tools/mock_qkd_nodes/start.sh &
-#    - sleep 10 # wait for nodes to spin up
-#  artifacts:
-#    paths:
-#      - mock_nodes.log
-#  rules:
-#    - if: '$CI_PIPELINE_SOURCE == "merge_request_event" && ($CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "develop" || $CI_MERGE_REQUEST_TARGET_BRANCH_NAME == $CI_DEFAULT_BRANCH)'
-#    - if: '$CI_PIPELINE_SOURCE == "push" && $CI_COMMIT_BRANCH == "develop"'
-
-## Prepare Scenario (Start NBI, mock services)
-#prepare_scenario:
-#  stage: deploy
-#  script:
-#    - pytest src/tests/qkd/unit/PrepareScenario.py
-#  needs:
-#    - start_mock_nodes
-#  rules:
-#    - if: '$CI_PIPELINE_SOURCE == "merge_request_event" && ($CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "develop" || $CI_MERGE_REQUEST_TARGET_BRANCH_NAME == $CI_DEFAULT_BRANCH)'
-#    - if: '$CI_PIPELINE_SOURCE == "push" && $CI_COMMIT_BRANCH == "develop"'
+# Deploy mock QKD nodes
+deploy_mock_qkd_nodes:
+  stage: prepare
+  script:
+    - echo "Starting mock QKD nodes deployment..."
+    - >
+      if docker network list | grep -q teraflowbridge; then
+        echo "teraflowbridge network is already created";
+      else
+        docker network create -d bridge teraflowbridge;
+      fi
+    - >
+      for port in 11111 22222 33333; do
+        if lsof -i:$port >/dev/null 2>&1; then
+          echo "Freeing up port $port...";
+          fuser -k $port/tcp;
+        fi
+      done
+    - MOCK_NODES_DIR="/home/cttc-docker/tfs-ctrl/src/tests/tools/mock_qkd_nodes"
+    - if [ -d "$MOCK_NODES_DIR" ]; then
+        cd "$MOCK_NODES_DIR" && ./start.sh &
+      else
+        echo "Error: Mock QKD nodes directory '$MOCK_NODES_DIR' not found." && exit 1;
+      fi
+  rules:
+    - if: '$CI_PIPELINE_SOURCE == "merge_request_event" && ($CI_MERGE_REQUEST_TARGET_BRANCH_NAME == "develop" || $CI_MERGE_REQUEST_TARGET_BRANCH_NAME == $CI_DEFAULT_BRANCH)'
+    - if: '$CI_PIPELINE_SOURCE == "push" && $CI_COMMIT_BRANCH == "develop"'
+    - changes:
+      - src/common/**/*.py
+      - proto/*.proto
+      - src/device/**/*.{py,in,yml}
+      - src/device/Dockerfile
+      - src/device/tests/*.py
+      - src/tests/tools/mock_qkd_nodes/**
+      - .gitlab-ci.yml
 
 # Apply unit test to the component
-unit_test device:
+unit_test_device:
   variables:
     IMAGE_NAME: 'device' # name of the microservice
     IMAGE_TAG: 'latest' # tag of the container image (production, development, etc)
   stage: unit_test
   needs:
-    - build device
-    #- start_mock_nodes
-    #- prepare_scenario
+    - build_device
+    - deploy_mock_qkd_nodes
   before_script:
     - docker login -u "$CI_REGISTRY_USER" -p "$CI_REGISTRY_PASSWORD" $CI_REGISTRY
-    - >
-      if docker network list | grep teraflowbridge; then
-        echo "teraflowbridge is already created";
-      else
-        docker network create -d bridge teraflowbridge;
-      fi
-    - >
-      if docker container ls | grep $IMAGE_NAME; then
-        docker rm -f $IMAGE_NAME;
-      else
-        echo "$IMAGE_NAME image is not in the system";
-      fi
   script:
     - docker pull "$CI_REGISTRY_IMAGE/$IMAGE_NAME:$IMAGE_TAG"
-    - docker run --name $IMAGE_NAME -d -p 2020:2020 -v "$PWD/src/$IMAGE_NAME/tests:/opt/results" --network=teraflowbridge $CI_REGISTRY_IMAGE/$IMAGE_NAME:$IMAGE_TAG
+    - docker run --name $IMAGE_NAME --network="host" -d -p 2020:2020 -v "$PWD/src/$IMAGE_NAME/tests:/opt/results" -e PYTHONPATH="/var/teraflow:/var/teraflow/device:/var/teraflow/tests/tools/mock_qkd_nodes:/var/teraflow/tests" "$CI_REGISTRY_IMAGE/$IMAGE_NAME:$IMAGE_TAG"
     - sleep 5
     - docker ps -a
     - docker logs $IMAGE_NAME
     - docker exec -i $IMAGE_NAME bash -c "coverage run --append -m pytest --log-level=INFO --verbose $IMAGE_NAME/tests/test_unitary_emulated.py --junitxml=/opt/results/${IMAGE_NAME}_report_emulated.xml"
     - docker exec -i $IMAGE_NAME bash -c "coverage run --append -m pytest --log-level=INFO --verbose $IMAGE_NAME/tests/test_unitary_ietf_actn.py --junitxml=/opt/results/${IMAGE_NAME}_report_ietf_actn.xml"
-    #- docker exec -i $IMAGE_NAME bash -c "coverage run --append -m pytest --log-level=INFO --verbose $IMAGE_NAME/tests/qkd/unit/test_*.py"
+    - docker exec -i $IMAGE_NAME bash -c "coverage run --append -m pytest --log-level=INFO --verbose device/tests/qkd/unit/test_qkd_mock_connectivity.py"
+    - docker exec -i $IMAGE_NAME bash -c "coverage run --append -m pytest --log-level=INFO --verbose device/tests/qkd/unit/test_qkd_compliance.py"
+    - docker exec -i $IMAGE_NAME bash -c "coverage run --append -m pytest --log-level=INFO --verbose device/tests/qkd/unit/test_mock_qkd_node.py"
+    - docker exec -i $IMAGE_NAME bash -c "coverage run --append -m pytest --log-level=INFO --verbose device/tests/qkd/unit/test_qkd_error_handling.py"
     - docker exec -i $IMAGE_NAME bash -c "coverage report --include='${IMAGE_NAME}/*' --show-missing"
   coverage: '/TOTAL\s+\d+\s+\d+\s+(\d+%)/'
   after_script:
@@ -110,12 +112,13 @@ unit_test device:
       - src/$IMAGE_NAME/Dockerfile
       - src/$IMAGE_NAME/tests/*.py
       - src/$IMAGE_NAME/tests/Dockerfile
+      - src/tests/tools/mock_qkd_nodes/**
       - manifests/${IMAGE_NAME}service.yaml
       - .gitlab-ci.yml
   artifacts:
-      when: always
-      reports:
-        junit: src/$IMAGE_NAME/tests/${IMAGE_NAME}_report_*.xml
+    when: always
+    reports:
+      junit: src/$IMAGE_NAME/tests/${IMAGE_NAME}_report_*.xml
 
 ## Deployment of the service in Kubernetes Cluster
 #deploy device:
diff --git a/src/device/tests/qkd/unit/test_mock_qkd_node.py b/src/device/tests/qkd/unit/test_mock_qkd_node.py
index f679a8389476c35a6881b7dc6484baab8ef4e20b..afd614fec8c07061977b8c0fdccf3ebfba71983b 100644
--- a/src/device/tests/qkd/unit/test_mock_qkd_node.py
+++ b/src/device/tests/qkd/unit/test_mock_qkd_node.py
@@ -16,16 +16,20 @@ import pytest
 import requests
 from requests.exceptions import ConnectionError
 
+MOCK_QKD_ADDRESS = '127.0.0.1'
+MOCK_PORT = 11111
+
+
 def test_mock_qkd_node_responses():
-    response = requests.get('http://127.0.0.1:11111/restconf/data/etsi-qkd-sdn-node:qkd_node')
+    response = requests.get(f'http://{MOCK_QKD_ADDRESS}:{MOCK_PORT}/restconf/data/etsi-qkd-sdn-node:qkd_node')
     assert response.status_code == 200
     data = response.json()
     assert 'qkd_node' in data
 
 def test_mock_node_failure_scenarios():
     try:
-        response = requests.get('http://127.0.0.1:12345/restconf/data/etsi-qkd-sdn-node:qkd_node')
+        response = requests.get(f'http://{MOCK_QKD_ADDRESS}:12345/restconf/data/etsi-qkd-sdn-node:qkd_node')
     except ConnectionError as e:
         assert isinstance(e, ConnectionError)
     else:
-        pytest.fail("ConnectionError not raised as expected")
+        pytest.fail("ConnectionError not raised as expected")
\ No newline at end of file
diff --git a/src/device/tests/qkd/unit/test_qkd_compliance.py b/src/device/tests/qkd/unit/test_qkd_compliance.py
index 2f305888e952e7d9acc3e96ffc1e427a7cc85685..d4e5cdaf65321eea6668b57bbfc193d57970b1d4 100644
--- a/src/device/tests/qkd/unit/test_qkd_compliance.py
+++ b/src/device/tests/qkd/unit/test_qkd_compliance.py
@@ -15,10 +15,18 @@
 
 import pytest
 import requests
+from requests.exceptions import HTTPError
 from tests.tools.mock_qkd_nodes.YangValidator import YangValidator
 
 def test_compliance_with_yang_models():
     validator = YangValidator('etsi-qkd-sdn-node', ['etsi-qkd-node-types'])
-    response = requests.get('http://127.0.0.1:11111/restconf/data/etsi-qkd-sdn-node:qkd_node')
-    data = response.json()
-    assert validator.parse_to_dict(data) is not None
+    try:
+        response = requests.get('http://127.0.0.1:11111/restconf/data/etsi-qkd-sdn-node:qkd_node')
+        response.raise_for_status()
+        data = response.json()
+        assert validator.parse_to_dict(data) is not None, "Data validation failed against YANG model."
+    except HTTPError as e:
+        pytest.fail(f"HTTP error occurred: {e}")
+    except Exception as e:
+        pytest.fail(f"Unexpected error occurred: {e}")
+
diff --git a/src/device/tests/qkd/unit/test_qkd_error_hanling.py b/src/device/tests/qkd/unit/test_qkd_error_handling.py
similarity index 96%
rename from src/device/tests/qkd/unit/test_qkd_error_hanling.py
rename to src/device/tests/qkd/unit/test_qkd_error_handling.py
index d93e3711136de496fd39365563032f827cfbe913..68f4532b3bd4f1fb92f6e6cad93f05e1bf5f6259 100644
--- a/src/device/tests/qkd/unit/test_qkd_error_hanling.py
+++ b/src/device/tests/qkd/unit/test_qkd_error_handling.py
@@ -40,7 +40,7 @@ def test_invalid_operations_on_network_links(qkd_driver):
 
     try:
         # Attempt to perform an invalid operation (simulate wrong resource key)
-        response = requests.post(f'http://{qkd_driver.address}/invalid_resource', json=invalid_payload)
+        response = requests.post(f'http://{qkd_driver.address}:{qkd_driver.port}/invalid_resource', json=invalid_payload)
         response.raise_for_status()
 
     except HTTPError as e:
diff --git a/src/device/tests/qkd/unit/test_qkd_mock_connectivity.py b/src/device/tests/qkd/unit/test_qkd_mock_connectivity.py
index be9427d9b619423a61a3e6f5270d8aab76dc8955..98a8b6d453aa0143c422eb0bfb5be4195221adde 100644
--- a/src/device/tests/qkd/unit/test_qkd_mock_connectivity.py
+++ b/src/device/tests/qkd/unit/test_qkd_mock_connectivity.py
@@ -12,16 +12,35 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest, requests
+import pytest
+import requests
+import time
+import socket
 from unittest.mock import patch
-from device.service.drivers.qkd.QKDDriver import QKDDriver
+from device.service.drivers.qkd.QKDDriver2 import QKDDriver
 
-MOCK_QKD_ADDRRESS = '127.0.0.1'
+MOCK_QKD_ADDRESS = '127.0.0.1'  # Use localhost to connect to the mock node in the Docker container
 MOCK_PORT = 11111
 
+@pytest.fixture(scope="module")
+def wait_for_mock_node():
+    """
+    Fixture to wait for the mock QKD node to be ready before running tests.
+    """
+    timeout = 30  # seconds
+    start_time = time.time()
+    while True:
+        try:
+            with socket.create_connection((MOCK_QKD_ADDRESS, MOCK_PORT), timeout=1):
+                break  # Success
+        except (socket.timeout, socket.error):
+            if time.time() - start_time > timeout:
+                raise RuntimeError("Timed out waiting for mock QKD node to be ready.")
+            time.sleep(1)
+
 @pytest.fixture
-def qkd_driver():
-    return QKDDriver(address=MOCK_QKD_ADDRRESS, port=MOCK_PORT, username='user', password='pass')
+def qkd_driver(wait_for_mock_node):
+    return QKDDriver(address=MOCK_QKD_ADDRESS, port=MOCK_PORT, username='user', password='pass')
 
 # Deliverable Test ID: SBI_Test_01
 def test_qkd_driver_connection(qkd_driver):
@@ -29,7 +48,7 @@ def test_qkd_driver_connection(qkd_driver):
 
 # Deliverable Test ID: SBI_Test_01
 def test_qkd_driver_invalid_connection():
-    qkd_driver = QKDDriver(address='127.0.0.1', port=12345, username='user', password='pass')  # Use invalid port directly
+    qkd_driver = QKDDriver(address=MOCK_QKD_ADDRESS, port=12345, username='user', password='pass')  # Use invalid port directly
     assert qkd_driver.Connect() is False
 
 # Deliverable Test ID: SBI_Test_10
@@ -38,4 +57,3 @@ def test_qkd_driver_timeout_connection(mock_get, qkd_driver):
     mock_get.side_effect = requests.exceptions.Timeout
     qkd_driver.timeout = 0.001  # Simulate very short timeout
     assert qkd_driver.Connect() is False
-
diff --git a/src/tests/tools/mock_qkd_nodes/wsgi.py b/src/tests/tools/mock_qkd_nodes/wsgi.py
index 3f8847849337fbfb1a9f84c783786218db4fb04d..b65cbc420a702789e0037fe911cdb14d1dbf2fbb 100644
--- a/src/tests/tools/mock_qkd_nodes/wsgi.py
+++ b/src/tests/tools/mock_qkd_nodes/wsgi.py
@@ -23,7 +23,7 @@ yang_validator = YangValidator('etsi-qkd-sdn-node', ['etsi-qkd-node-types'])
 
 
 nodes = {
-    '10.0.2.10:11111': {'node': {
+    '127.0.0.1:11111': {'node': {
             'qkdn_id': '00000001-0000-0000-0000-000000000000',
         },
         'qkdn_capabilities': {
@@ -54,7 +54,7 @@ nodes = {
                 {
                     'qkdi_id': '101',
                     'qkdi_att_point': {
-                        'device':'10.0.2.10',
+                        'device':'127.0.0.1',
                         'port':'1001'
                     },
                     'qkdi_capabilities': {
@@ -69,7 +69,7 @@ nodes = {
         }
     },
 
-    '10.0.2.10:22222': {'node': {
+    '127.0.0.1:22222': {'node': {
             'qkdn_id': '00000002-0000-0000-0000-000000000000',
         },
         'qkdn_capabilities': {
@@ -100,7 +100,7 @@ nodes = {
                 {
                     'qkdi_id': '201',
                     'qkdi_att_point': {
-                        'device':'10.0.2.10',
+                        'device':'127.0.0.1',
                         'port':'2001'
                     },
                     'qkdi_capabilities': {
@@ -109,7 +109,7 @@ nodes = {
                 {
                     'qkdi_id': '202',
                     'qkdi_att_point': {
-                        'device':'10.0.2.10',
+                        'device':'127.0.0.1',
                         'port':'2002'
                     },
                     'qkdi_capabilities': {
@@ -124,7 +124,7 @@ nodes = {
         }
     },
 
-    '10.0.2.10:33333': {'node': {
+    '127.0.0.1:33333': {'node': {
             'qkdn_id': '00000003-0000-0000-0000-000000000000',
         },
         'qkdn_capabilities': {
@@ -155,7 +155,7 @@ nodes = {
                 {
                     'qkdi_id': '301',
                     'qkdi_att_point': {
-                        'device':'10.0.2.10',
+                        'device':'127.0.0.1',
                         'port':'3001'
                     },
                     'qkdi_capabilities': {
@@ -316,53 +316,3 @@ def patch(path):
         reason = str(e)
         success = False
     return {'success': success, 'reason': reason}
-
-
-
-
-
-# import json
-# from mock import requests
-# import pyangbind.lib.pybindJSON as enc
-# from pyangbind.lib.serialise import pybindJSONDecoder as dec
-# from yang.sbi.qkd.templates.etsi_qkd_sdn_node import etsi_qkd_sdn_node
-
-# module = etsi_qkd_sdn_node()
-# url = 'https://1.1.1.1/restconf/data/etsi-qkd-sdn-node:'
-
-# # Get node all info
-# z = requests.get(url).json()
-# var = dec.load_json(z, None, None, obj=module)
-# print(enc.dumps(var))
-
-
-# Reset module variable because it is already filled
-# module = etsi_qkd_sdn_node()
-
-# # Get node basic info
-# node = module.qkd_node
-# z = requests.get(url + 'qkd_node').json()
-# var = dec.load_json(z, None, None, obj=node)
-# print(enc.dumps(var))
-
-
-# # Get all apps
-# apps = node.qkd_applications
-# z = requests.get(url + 'qkd_node/qkd_applications').json()
-# var = dec.load_json(z, None, None, obj=apps)
-# print(enc.dumps(var))
-
-# # Edit app 0
-# app = apps.qkd_app['00000000-0001-0000-0000-000000000000']
-# app.client_app_id = 'id_0'
-# requests.put(url + 'qkd_node/qkd_applications/qkd_app=00000000-0001-0000-0000-000000000000', json=json.loads(enc.dumps(app)))
-
-# # Create app 1
-# app = apps.qkd_app.add('00000000-0001-0000-0000-000000000001')
-# requests.post(url + 'qkd_node/qkd_applications/qkd_app=00000000-0001-0000-0000-000000000001', json=json.loads(enc.dumps(app)))
-
-# # Get all apps
-# apps = node.qkd_applications
-# z = requests.get(url + 'qkd_node/qkd_applications').json()
-# var = dec.load_json(z, None, None, obj=apps)
-# print(enc.dumps(var))