import io

import pandas as pd
import pytest


@pytest.fixture
def minimal_df() -> pd.DataFrame:
    return pd.DataFrame(
        {"protocol": ["TCP", "UDP", "TCP", "UDP"],
         "label": ["benign", "malicious", "benign", "malicious"]}
    )


@pytest.fixture
def csv_buffer_two_classes() -> io.StringIO:
    buf = io.StringIO("protocol,label\nTCP,benign\nUDP,malicious\nTCP,benign\n")
    buf.seek(0)
    return buf
