Cybersecurity-Projects/PROJECTS/beginner/network-traffic-analyzer/python/tests/test_models.py

"""
ⒸAngelaMos | 2026
test_models.py

Basic happy path tests for data models

Covers Protocol enum values, PacketInfo field storage and optional field
defaults, CaptureConfig defaults and custom values, EndpointStats computed
totals, and CaptureStatistics computed properties including duration,
average bandwidth, protocol percentages, and top talkers sorting.

Tests:
  TestProtocol - all Protocol enum string values
  TestPacketInfo - field storage and optional field defaults
  TestCaptureConfig - default and custom configuration values
  TestEndpointStats - total_packets and total_bytes computed properties
  TestCaptureStatistics - empty state, duration, bandwidth, percentages, top talkers

Connects to:
  models.py - all symbols imported and tested here
"""


import pytest

from netanal.models import (
    CaptureConfig,
    CaptureStatistics,
    EndpointStats,
    PacketInfo,
    Protocol,
)


class TestProtocol:
    """
    Tests for the Protocol enum
    """
    def test_protocol_values(self):
        """
        Verify all protocol enum values match expected strings
        """
        assert Protocol.TCP.value == "TCP"
        assert Protocol.UDP.value == "UDP"
        assert Protocol.ICMP.value == "ICMP"
        assert Protocol.DNS.value == "DNS"
        assert Protocol.HTTP.value == "HTTP"
        assert Protocol.HTTPS.value == "HTTPS"
        assert Protocol.ARP.value == "ARP"
        assert Protocol.OTHER.value == "OTHER"


class TestPacketInfo:
    """
    Tests for the PacketInfo dataclass
    """
    def test_create_packet_info(self):
        """
        Verify PacketInfo stores all fields correctly
        """
        packet = PacketInfo(
            timestamp=1234567890.123,
            src_ip="192.168.1.1",
            dst_ip="192.168.1.2",
            protocol=Protocol.TCP,
            size=1500,
            src_port=443,
            dst_port=54321,
        )

        assert packet.timestamp == pytest.approx(1234567890.123)
        assert packet.src_ip == "192.168.1.1"
        assert packet.dst_ip == "192.168.1.2"
        assert packet.protocol == Protocol.TCP
        assert packet.size == 1500
        assert packet.src_port == 443
        assert packet.dst_port == 54321

    def test_packet_info_optional_fields(self):
        """
        Verify optional fields default to None
        """
        packet = PacketInfo(
            timestamp=0.0,
            src_ip="10.0.0.1",
            dst_ip="10.0.0.2",
            protocol=Protocol.ICMP,
            size=64,
        )

        assert packet.src_port is None
        assert packet.dst_port is None
        assert packet.src_mac is None
        assert packet.dst_mac is None


class TestCaptureConfig:
    """
    Tests for the CaptureConfig dataclass
    """
    def test_default_config(self):
        """
        Verify default configuration values
        """
        config = CaptureConfig()

        assert config.interface is None
        assert config.bpf_filter is None
        assert config.packet_count is None
        assert config.timeout_seconds is None
        assert config.store_packets is False

    def test_custom_config(self):
        """
        Verify custom configuration is stored correctly
        """
        config = CaptureConfig(
            interface="eth0",
            bpf_filter="tcp port 80",
            packet_count=100,
            timeout_seconds=30.0,
        )

        assert config.interface == "eth0"
        assert config.bpf_filter == "tcp port 80"
        assert config.packet_count == 100
        assert config.timeout_seconds == pytest.approx(30.0)


class TestEndpointStats:
    """
    Tests for the EndpointStats dataclass
    """
    def test_endpoint_stats_totals(self):
        """
        Verify total_packets and total_bytes computed properties
        """
        endpoint = EndpointStats(ip_address="192.168.1.100")
        endpoint.packets_sent = 50
        endpoint.packets_received = 30
        endpoint.bytes_sent = 5000
        endpoint.bytes_received = 3000

        assert endpoint.total_packets == 80
        assert endpoint.total_bytes == 8000


class TestCaptureStatistics:
    """
    Tests for the CaptureStatistics dataclass
    """
    def test_empty_statistics(self):
        """
        Verify empty statistics have zero values
        """
        stats = CaptureStatistics()

        assert stats.total_packets == 0
        assert stats.total_bytes == 0
        assert len(stats.protocol_distribution) == 0
        assert len(stats.endpoints) == 0

    def test_duration_calculation(self):
        """
        Verify duration_seconds computed property
        """
        stats = CaptureStatistics(
            start_time=1000.0,
            end_time=1010.0,
        )

        assert stats.duration_seconds == pytest.approx(10.0)

    def test_average_bandwidth(self):
        """
        Verify average_bandwidth calculation (bytes/second)
        """
        stats = CaptureStatistics(
            start_time=1000.0,
            end_time=1010.0,
            total_bytes=10000,
        )

        assert stats.average_bandwidth == pytest.approx(1000.0)

    def test_protocol_percentages(self):
        """
        Verify get_protocol_percentages returns correct distribution
        """
        stats = CaptureStatistics(total_packets=100)
        stats.protocol_distribution[Protocol.TCP] = 70
        stats.protocol_distribution[Protocol.UDP] = 30

        percentages = stats.get_protocol_percentages()

        assert percentages[Protocol.TCP] == pytest.approx(70.0)
        assert percentages[Protocol.UDP] == pytest.approx(30.0)

    def test_top_talkers(self):
        """
        Verify get_top_talkers returns endpoints sorted by total bytes
        """
        stats = CaptureStatistics()

        endpoint1 = EndpointStats(ip_address="192.168.1.1")
        endpoint1.bytes_sent = 1000
        endpoint1.bytes_received = 500

        endpoint2 = EndpointStats(ip_address="192.168.1.2")
        endpoint2.bytes_sent = 5000
        endpoint2.bytes_received = 2000

        endpoint3 = EndpointStats(ip_address="192.168.1.3")
        endpoint3.bytes_sent = 100
        endpoint3.bytes_received = 50

        stats.endpoints["192.168.1.1"] = endpoint1
        stats.endpoints["192.168.1.2"] = endpoint2
        stats.endpoints["192.168.1.3"] = endpoint3

        top = stats.get_top_talkers(limit=2)

        assert len(top) == 2
        assert top[0].ip_address == "192.168.1.2"
        assert top[1].ip_address == "192.168.1.1"