Include envorunment variables for qkds

import pytest

import json

import logging

from src.device.service.drivers.qkd.QKDDriver2 import QKDDriver

import os

from dotenv import load_dotenv

from src.device.service.drivers.qkd.QKDDriver import QKDDriver

# Load environment variables from .env file

load_dotenv()

# Set up logging

logging.basicConfig(level=logging.INFO)

    "state_qkd2": None,

}

# Environment variables for sensitive information

QKD1_ADDRESS = os.getenv("QKD1_ADDRESS")

QKD2_ADDRESS = os.getenv("QKD2_ADDRESS")

PORT = os.getenv("QKD_PORT")

USERNAME = os.getenv("QKD_USERNAME")

PASSWORD = os.getenv("QKD_PASSWORD")

@pytest.fixture

def driver_qkd1():

    return QKDDriver(address='10.13.13.2', port=5100, username='admin', password='password', use_jwt=True)

    return QKDDriver(address=QKD1_ADDRESS, port=PORT, username=USERNAME, password=PASSWORD, use_jwt=True)

@pytest.fixture

def driver_qkd2():

    return QKDDriver(address='10.13.13.3', port=5100, username='admin', password='password', use_jwt=True)

    return QKDDriver(address=QKD2_ADDRESS, port=PORT, username=USERNAME, password=PASSWORD, use_jwt=True)

def log_data(label, data):

    """Logs data in JSON format with a label."""

    LOGGER.info(f"{label}: {json.dumps(data, indent=2, cls=SafeJSONEncoder)}")

def get_jwt_token(driver):

    """Retrieve JWT token from the driver."""

    try:

        return driver._QKDDriver__headers.get('Authorization').split(' ')[1]

    except (AttributeError, KeyError, TypeError):

        LOGGER.error("Failed to retrieve JWT token")

        return None

def save_json_file(filename, data):

    except Exception as e:

        LOGGER.error(f"Failed to save {filename}: {e}")

def retrieve_data(driver, label, method, *args):

    """Retrieve data from the driver and log it."""

    try:

        data = method(*args)

        log_data(label, data)

        return data

    except Exception as e:

        LOGGER.error(f"Failed to retrieve {label}: {e}")

        return None

def test_retrieve_and_create_descriptor(driver_qkd1, driver_qkd2):

    # Connect to both QKD nodes

    assert driver_qkd1.Connect()

    assert driver_qkd2.Connect()

    assert driver_qkd1.Connect(), "Failed to connect to QKD1"

    assert driver_qkd2.Connect(), "Failed to connect to QKD2"

    # Use the same JWT token for all requests

    jwt_token = get_jwt_token(driver_qkd1)

    assert jwt_token, "Failed to retrieve JWT token from QKD1"

    driver_qkd2._QKDDriver__headers['Authorization'] = f'Bearer {jwt_token}'

    # Retrieve initial configs

    config_qkd1 = driver_qkd1.GetInitialConfig()

    retrieved_info['config_qkd1'] = config_qkd1

    log_data("QKD1 Initial Config", config_qkd1)

    assert config_qkd1, "Failed to retrieve initial configuration for QKD1"

    config_qkd2 = driver_qkd2.GetInitialConfig()

    retrieved_info['config_qkd2'] = config_qkd2

    log_data("QKD2 Initial Config", config_qkd2)

    assert config_qkd2, "Failed to retrieve initial configuration for QKD2"

    # Retrieve configurations

    retrieved_info['config_qkd1'] = retrieve_data(driver_qkd1, "QKD1 Initial Config", driver_qkd1.GetInitialConfig)

    retrieved_info['config_qkd2'] = retrieve_data(driver_qkd2, "QKD2 Initial Config", driver_qkd2.GetInitialConfig)

    # Retrieve capabilities

    capabilities_qkd1 = driver_qkd1.GetConfig(['capabilities'])

    retrieved_info['capabilities_qkd1'] = capabilities_qkd1

    log_data("QKD1 Capabilities", capabilities_qkd1)

    assert capabilities_qkd1, "Failed to retrieve capabilities for QKD1"

    capabilities_qkd2 = driver_qkd2.GetConfig(['capabilities'])

    retrieved_info['capabilities_qkd2'] = capabilities_qkd2

    log_data("QKD2 Capabilities", capabilities_qkd2)

    assert capabilities_qkd2, "Failed to retrieve capabilities for QKD2"

    retrieved_info['capabilities_qkd1'] = retrieve_data(driver_qkd1, "QKD1 Capabilities", driver_qkd1.GetConfig, ['capabilities'])

    retrieved_info['capabilities_qkd2'] = retrieve_data(driver_qkd2, "QKD2 Capabilities", driver_qkd2.GetConfig, ['capabilities'])

    # Retrieve interfaces

    interfaces_qkd1 = driver_qkd1.GetConfig(['interfaces'])

    retrieved_info['interfaces_qkd1'] = interfaces_qkd1

    log_data("QKD1 Interfaces", interfaces_qkd1)

    assert interfaces_qkd1, "Failed to retrieve interfaces for QKD1"

    interfaces_qkd2 = driver_qkd2.GetConfig(['interfaces'])

    retrieved_info['interfaces_qkd2'] = interfaces_qkd2

    log_data("QKD2 Interfaces", interfaces_qkd2)

    assert interfaces_qkd2, "Failed to retrieve interfaces for QKD2"

    retrieved_info['interfaces_qkd1'] = retrieve_data(driver_qkd1, "QKD1 Interfaces", driver_qkd1.GetConfig, ['interfaces'])

    retrieved_info['interfaces_qkd2'] = retrieve_data(driver_qkd2, "QKD2 Interfaces", driver_qkd2.GetConfig, ['interfaces'])

    # Retrieve links

    links_qkd1 = driver_qkd1.GetConfig(['links'])

    retrieved_info['links_qkd1'] = links_qkd1

    log_data("QKD1 Links", links_qkd1)

    assert links_qkd1, "Failed to retrieve links for QKD1"

    links_qkd2 = driver_qkd2.GetConfig(['links'])

    retrieved_info['links_qkd2'] = links_qkd2

    log_data("QKD2 Links", links_qkd2)

    assert links_qkd2, "Failed to retrieve links for QKD2"

    retrieved_info['links_qkd1'] = retrieve_data(driver_qkd1, "QKD1 Links", driver_qkd1.GetConfig, ['links'])

    retrieved_info['links_qkd2'] = retrieve_data(driver_qkd2, "QKD2 Links", driver_qkd2.GetConfig, ['links'])

    # Retrieve states

    state_qkd1 = driver_qkd1.GetState()

    retrieved_info['state_qkd1'] = state_qkd1

    log_data("QKD1 Current State", state_qkd1)

    assert state_qkd1, "Failed to retrieve state for QKD1"

    state_qkd2 = driver_qkd2.GetState()

    retrieved_info['state_qkd2'] = state_qkd2

    log_data("QKD2 Current State", state_qkd2)

    assert state_qkd2, "Failed to retrieve state for QKD2"

    retrieved_info['state_qkd1'] = retrieve_data(driver_qkd1, "QKD1 Current State", driver_qkd1.GetState)

    retrieved_info['state_qkd2'] = retrieve_data(driver_qkd2, "QKD2 Current State", driver_qkd2.GetState)

    # Save retrieved information after all data retrieval

    # Save retrieved information

    save_json_file('retrieved_info.json', retrieved_info)

    # Dynamically create the descriptor

    # Create descriptor dynamically

    descriptor = {

        "contexts": [

            {"context_id": {"context_uuid": {"uuid": "admin"}}}

        ],

        "topologies": [

            {"topology_id": {"topology_uuid": {"uuid": "admin"}, "context_id": {"context_uuid": {"uuid": "admin"}}}}

        ],

        "contexts": [{"context_id": {"context_uuid": {"uuid": "admin"}}}],

        "topologies": [{"topology_id": {"topology_uuid": {"uuid": "admin"}, "context_id": {"context_uuid": {"uuid": "admin"}}}}],

        "devices": [],

        "links": []

    }

    # Dynamically add device information

    for config, token, interfaces, device_name, address, port in [

        (config_qkd1, jwt_token, interfaces_qkd1, "QKD1", "10.13.13.2", "5100"),

        (config_qkd2, jwt_token, interfaces_qkd2, "QKD2", "10.13.13.3", "5100")

    # Add device information to descriptor

    for config, token, interfaces, device_name, address in [

        (retrieved_info['config_qkd1'], jwt_token, retrieved_info['interfaces_qkd1'], "QKD1", QKD1_ADDRESS),

        (retrieved_info['config_qkd2'], jwt_token, retrieved_info['interfaces_qkd2'], "QKD2", QKD2_ADDRESS)

    ]:

        device_info = {

            "device_id": {"device_uuid": {"uuid": device_name}},

            "device_type": "qkd-node",

            "device_operational_status": 0,

            "device_drivers": [12],  # This could be dynamically determined if needed

            "device_drivers": [12],

            "device_endpoints": [],

            "device_config": {

                "config_rules": [

                    {"action": 1, "custom": {"resource_key": "_connect/address", "resource_value": address}},

                    {"action": 1, "custom": {"resource_key": "_connect/port", "resource_value": port}},

                    {"action": 1, "custom": {"resource_key": "_connect/settings", "resource_value": {

                        "scheme": "http",

                        "token": token if token else "N/A"

                    }}}

                    {"action": 1, "custom": {"resource_key": "_connect/port", "resource_value": PORT}},

                    {"action": 1, "custom": {"resource_key": "_connect/settings", "resource_value": {"scheme": "http", "token": token}}}

                ]

            }

        }

        descriptor['devices'].append(device_info)

    # Dynamically create and add links based on retrieved links data

    if links_qkd1 and links_qkd2:

        for link_data in links_qkd1:

            link_uuid = link_data[1].get('qkdl_id')

    # Create links based on retrieved link data

    if retrieved_info['links_qkd1'] and retrieved_info['links_qkd2']:

        for link_data in retrieved_info['links_qkd1']:

            link_entry = {

                "link_id": {"link_uuid": {"uuid": f"QKD1/{address}:{port}==QKD2/{links_qkd2[0][1]['qkdi_status']}/{port}"}},

                "link_id": {"link_uuid": {"uuid": f"QKD1/{QKD1_ADDRESS}:{PORT}==QKD2/{retrieved_info['links_qkd2'][0][1]['qkdi_status']}/{PORT}"}},

                "link_endpoint_ids": [

                    {"device_id": {"device_uuid": {"uuid": "QKD1"}}, "endpoint_uuid": {"uuid": f"{address}:{port}"}},

                    {"device_id": {"device_uuid": {"uuid": "QKD2"}}, "endpoint_uuid": {"uuid": f"{address}:{port}"}}

                    {"device_id": {"device_uuid": {"uuid": "QKD1"}}, "endpoint_uuid": {"uuid": f"{QKD1_ADDRESS}:{PORT}"}},

                    {"device_id": {"device_uuid": {"uuid": "QKD2"}}, "endpoint_uuid": {"uuid": f"{QKD2_ADDRESS}:{PORT}"}}

                ]

            }

            descriptor['links'].append(link_entry)