Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: OBS-441 - reconciler: use manufacturer provided with platform + improvements #80

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
212 changes: 126 additions & 86 deletions diode-sdk-python/netboxlabs/diode/sdk/ingester.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
# Copyright 2024 NetBox Labs Inc
"""NetBox Labs, Diode - SDK - ingester protobuf message wrappers."""

from typing import Mapping as _Mapping
from typing import Any
from typing import Optional as _Optional
from typing import Union as _Union

Expand Down Expand Up @@ -43,6 +43,13 @@
)


def convert_to_protobuf(value: Any, protobuf_class, **kwargs):
"""Convert a value to a protobuf message."""
if isinstance(value, str):
return protobuf_class(**kwargs)
return value


class Manufacturer:
"""Manufacturer message wrapper."""

Expand Down Expand Up @@ -77,8 +84,9 @@ def __new__(
tags: _Optional[list[str]] = None,
) -> PlatformPb:
"""Create a new Platform protobuf message."""
if isinstance(manufacturer, str):
manufacturer = ManufacturerPb(name=manufacturer)
manufacturer = convert_to_protobuf(
manufacturer, ManufacturerPb, name=manufacturer
)

if isinstance(tags, list) and all(isinstance(t, str) for t in tags):
tags = [TagPb(name=tag) for tag in tags]
Expand Down Expand Up @@ -130,8 +138,9 @@ def __new__(
tags: _Optional[list[str]] = None,
) -> DeviceTypePb:
"""Create a new DeviceType protobuf message."""
if isinstance(manufacturer, str):
manufacturer = ManufacturerPb(name=manufacturer)
manufacturer = convert_to_protobuf(
manufacturer, ManufacturerPb, name=manufacturer
)

if isinstance(tags, list) and all(isinstance(t, str) for t in tags):
tags = [TagPb(name=tag) for tag in tags]
Expand Down Expand Up @@ -169,29 +178,39 @@ def __new__(
manufacturer: _Optional[_Union[str, Manufacturer, ManufacturerPb]] = None,
) -> DevicePb:
"""Create a new Device protobuf message."""
if isinstance(manufacturer, str):
manufacturer = ManufacturerPb(name=manufacturer)
manufacturer = convert_to_protobuf(
manufacturer, ManufacturerPb, name=manufacturer
)
platform = convert_to_protobuf(
platform, PlatformPb, name=platform, manufacturer=manufacturer
)

if isinstance(platform, str):
platform = PlatformPb(name=platform, manufacturer=manufacturer)
if (
isinstance(platform, PlatformPb)
and not platform.HasField("manufacturer")
and manufacturer is not None
):
platform.manufacturer.CopyFrom(manufacturer)

if isinstance(site, str):
site = SitePb(name=site)
site = convert_to_protobuf(site, SitePb, name=site)
device_type = convert_to_protobuf(
device_type, DeviceTypePb, model=device_type, manufacturer=manufacturer
)

if isinstance(device_type, str):
device_type = DeviceTypePb(model=device_type, manufacturer=manufacturer)
if (
isinstance(device_type, DeviceTypePb)
and not device_type.HasField("manufacturer")
and manufacturer is not None
):
device_type.manufacturer.CopyFrom(manufacturer)

if isinstance(role, str):
role = RolePb(name=role)
role = convert_to_protobuf(role, RolePb, name=role)

if isinstance(tags, list) and all(isinstance(t, str) for t in tags):
tags = [TagPb(name=tag) for tag in tags]

if isinstance(primary_ip4, str):
primary_ip4 = IPAddressPb(address=primary_ip4)

if isinstance(primary_ip6, str):
primary_ip6 = IPAddressPb(address=primary_ip6)
primary_ip4 = convert_to_protobuf(primary_ip4, IPAddressPb, address=IPAddressPb)
primary_ip6 = convert_to_protobuf(primary_ip6, IPAddressPb, address=IPAddressPb)

return DevicePb(
name=name,
Expand Down Expand Up @@ -236,29 +255,45 @@ def __new__(
tags: _Optional[list[str]] = None,
) -> InterfacePb:
"""Create a new Interface protobuf message."""
if isinstance(manufacturer, str):
manufacturer = ManufacturerPb(name=manufacturer)
manufacturer = convert_to_protobuf(
manufacturer, ManufacturerPb, name=manufacturer
)

platform = convert_to_protobuf(
platform, PlatformPb, name=platform, manufacturer=manufacturer
)

if (
isinstance(platform, PlatformPb)
and not platform.HasField("manufacturer")
and manufacturer is not None
):
platform.manufacturer.CopyFrom(manufacturer)

if isinstance(platform, str):
platform = PlatformPb(name=platform, manufacturer=manufacturer)
site = convert_to_protobuf(site, SitePb, name=site)

if isinstance(site, str):
site = SitePb(name=site)
device_type = convert_to_protobuf(
device_type, DeviceTypePb, model=device_type, manufacturer=manufacturer
)

if isinstance(device_type, str):
device_type = DeviceTypePb(model=device_type, manufacturer=manufacturer)
if (
isinstance(device_type, DeviceTypePb)
and not device_type.HasField("manufacturer")
and manufacturer is not None
):
device_type.manufacturer.CopyFrom(manufacturer)

if isinstance(role, str):
role = RolePb(name=role)
role = convert_to_protobuf(role, RolePb, name=role)

if isinstance(device, str):
device = DevicePb(
name=device,
device_type=device_type,
platform=platform,
site=site,
role=role,
)
device = convert_to_protobuf(
device,
DevicePb,
name=device,
device_type=device_type,
platform=platform,
site=site,
role=role,
)

if isinstance(tags, list) and all(isinstance(t, str) for t in tags):
tags = [TagPb(name=tag) for tag in tags]
Expand Down Expand Up @@ -301,31 +336,52 @@ def __new__(
tags: _Optional[list[str]] = None,
) -> IPAddressPb:
"""Create a new IPAddress protobuf message."""
if isinstance(manufacturer, str):
manufacturer = ManufacturerPb(name=manufacturer)
manufacturer = convert_to_protobuf(
manufacturer, ManufacturerPb, name=manufacturer
)

if isinstance(platform, str):
platform = PlatformPb(name=platform, manufacturer=manufacturer)
platform = convert_to_protobuf(
platform, PlatformPb, name=platform, manufacturer=manufacturer
)

if isinstance(site, str):
site = SitePb(name=site)
if (
isinstance(platform, PlatformPb)
and not platform.HasField("manufacturer")
and manufacturer is not None
):
platform.manufacturer.CopyFrom(manufacturer)

if isinstance(device_type, str):
device_type = DeviceTypePb(model=device_type, manufacturer=manufacturer)
site = convert_to_protobuf(site, SitePb, name=site)

if isinstance(device_role, str):
device_role = RolePb(name=device_role)
device_type = convert_to_protobuf(
device_type, DeviceTypePb, model=device_type, manufacturer=manufacturer
)

if (
isinstance(device_type, DeviceTypePb)
and not device_type.HasField("manufacturer")
and manufacturer is not None
):
device_type.manufacturer.CopyFrom(manufacturer)

device_role = convert_to_protobuf(device_role, RolePb, name=device_role)

device = convert_to_protobuf(
device,
DevicePb,
name=device,
device_type=device_type,
platform=platform,
site=site,
role=device_role,
)

if isinstance(device, str):
device = DevicePb(
name=device,
device_type=device_type,
platform=platform,
site=site,
role=device_role,
)
if isinstance(interface, str):
interface = InterfacePb(name=interface, device=device)
interface = convert_to_protobuf(
interface,
InterfacePb,
name=interface,
device=device,
)

if isinstance(tags, list) and all(isinstance(t, str) for t in tags):
tags = [TagPb(name=tag) for tag in tags]
Expand Down Expand Up @@ -357,8 +413,7 @@ def __new__(
tags: _Optional[list[str]] = None,
) -> PrefixPb:
"""Create a new Prefix protobuf message."""
if isinstance(site, str):
site = SitePb(name=site)
site = convert_to_protobuf(site, SitePb, name=site)

if isinstance(tags, list) and all(isinstance(t, str) for t in tags):
tags = [TagPb(name=tag) for tag in tags]
Expand Down Expand Up @@ -422,32 +477,17 @@ def __new__(
timestamp: _Optional[_timestamp_pb2.Timestamp] = None,
):
"""Create a new Entity protobuf message."""
if isinstance(site, str):
site = SitePb(name=site)

if isinstance(platform, str):
platform = PlatformPb(name=platform)

if isinstance(manufacturer, str):
manufacturer = ManufacturerPb(name=manufacturer)

if isinstance(device, str):
device = DevicePb(name=device)

if isinstance(device_role, str):
device_role = RolePb(name=device_role)

if isinstance(device_type, str):
device_type = DeviceTypePb(model=device_type)

if isinstance(ip_address, str):
ip_address = IPAddressPb(address=ip_address)

if isinstance(interface, str):
interface = InterfacePb(name=interface)

if isinstance(prefix, str):
prefix = PrefixPb(prefix=prefix)
site = convert_to_protobuf(site, SitePb, name=site)
platform = convert_to_protobuf(platform, PlatformPb, name=platform)
manufacturer = convert_to_protobuf(
manufacturer, ManufacturerPb, name=manufacturer
)
device = convert_to_protobuf(device, DevicePb, name=device)
device_role = convert_to_protobuf(device_role, RolePb, name=device_role)
device_type = convert_to_protobuf(device_type, DeviceTypePb, model=device_type)
ip_address = convert_to_protobuf(ip_address, IPAddressPb, address=ip_address)
interface = convert_to_protobuf(interface, InterfacePb, name=interface)
prefix = convert_to_protobuf(prefix, PrefixPb, prefix=prefix)

return EntityPb(
site=site,
Expand Down
68 changes: 68 additions & 0 deletions diode-sdk-python/tests/test_wrappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,24 @@ def test_device_wrapper():
assert device.asset_tag == "123456"
assert device.status == "active"

device = Device(
name="Device ABC",
device_type=DeviceType(model="Device Type ABC"),
platform=Platform(name="Platform ABC"),
manufacturer="Cisco",
)
assert device.device_type.manufacturer.name == "Cisco"
assert device.platform.manufacturer.name == "Cisco"

device = Device(
name="Device ABC",
device_type=DeviceType(model="Device Type ABC", manufacturer="Manufacturer A"),
platform=Platform(name="Platform ABC", manufacturer="Manufacturer A"),
manufacturer="Cisco",
)
assert device.device_type.manufacturer.name == "Manufacturer A"
assert device.platform.manufacturer.name == "Manufacturer A"


def test_interface_wrapper():
"""Ensure Interface wrapper instantiates InterfacePb."""
Expand Down Expand Up @@ -254,6 +272,30 @@ def test_interface_wrapper():
assert interface.mac_address == "00:00:00:00:00:00"
assert interface.description == "Description ABC"

interface = Interface(
name="Interface ABC",
device="Device ABC",
device_type="Device Type ABC",
role="Role ABC",
platform="Platform ABC",
site="Site ABC",
manufacturer="Cisco",
)
assert interface.device.device_type.manufacturer.name == "Cisco"
assert interface.device.platform.manufacturer.name == "Cisco"

interface = Interface(
name="Interface ABC",
device="Device ABC",
device_type=DeviceType(model="Device Type ABC", manufacturer="Manufacturer A"),
role="Role ABC",
platform=Platform(name="Platform ABC", manufacturer="Manufacturer A"),
site="Site ABC",
manufacturer="Cisco",
)
assert interface.device.device_type.manufacturer.name == "Manufacturer A"
assert interface.device.platform.manufacturer.name == "Manufacturer A"


def test_ip_address_wrapper():
"""Ensure IPAddress wrapper instantiates IPAddressPb."""
Expand Down Expand Up @@ -301,6 +343,32 @@ def test_ip_address_wrapper():
assert isinstance(ip_address.interface.device.platform.manufacturer, ManufacturerPb)
assert ip_address.interface.device.platform.manufacturer.name == "Cisco"

ip_address = IPAddress(
address="192.168.0.1/24",
interface="Interface ABC",
device="Device ABC",
device_type="Device Type ABC",
device_role="Role ABC",
platform="Platform ABC",
manufacturer="Cisco",
site="Site ABC",
)
assert ip_address.interface.device.device_type.manufacturer.name == "Cisco"
assert ip_address.interface.device.platform.manufacturer.name == "Cisco"

ip_address = IPAddress(
address="192.168.0.1/24",
interface="Interface ABC",
device="Device ABC",
device_type=DeviceType(model="Device Type ABC", manufacturer="Manufacturer A"),
device_role="Role ABC",
platform=Platform(name="Platform ABC", manufacturer="Manufacturer A"),
manufacturer="Cisco",
site="Site ABC",
)
assert ip_address.interface.device.device_type.manufacturer.name == "Manufacturer A"
assert ip_address.interface.device.platform.manufacturer.name == "Manufacturer A"


def test_prefix_wrapper():
"""Ensure Prefix wrapper instantiates PrefixPb."""
Expand Down
Loading
Loading