derp/tests/test_dns_plugin.py

"""Tests for the DNS lookup plugin."""

import asyncio
import importlib.util
import sys
from pathlib import Path
from unittest.mock import AsyncMock, patch

from derp.irc import Message

# plugins/ is not a Python package -- load the module from file path
_spec = importlib.util.spec_from_file_location(
    "plugins.dns", Path(__file__).resolve().parent.parent / "plugins" / "dns.py",
)
_mod = importlib.util.module_from_spec(_spec)
sys.modules[_spec.name] = _mod
_spec.loader.exec_module(_mod)

from plugins.dns import cmd_dns  # noqa: E402

# -- Helpers -----------------------------------------------------------------

class _FakeBot:
    def __init__(self):
        self.replied: list[str] = []

    async def reply(self, message, text: str) -> None:
        self.replied.append(text)


def _msg(text: str) -> Message:
    return Message(
        raw="", prefix="alice!~alice@host", nick="alice",
        command="PRIVMSG", params=["#test", text], tags={},
    )


def _mock_query(rcode, results):
    """Create an AsyncMock returning (rcode, results)."""
    mock = AsyncMock(return_value=(rcode, results))
    return mock


# -- Auto-detect type --------------------------------------------------------

class TestAutoDetect:
    def test_domain_defaults_to_a(self):
        bot = _FakeBot()
        mock = _mock_query(0, ["93.184.216.34"])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com")))
        assert "A:" in bot.replied[0]
        assert "93.184.216.34" in bot.replied[0]

    def test_ip_defaults_to_ptr(self):
        bot = _FakeBot()
        mock = _mock_query(0, ["dns.google"])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns 8.8.8.8")))
        assert "PTR:" in bot.replied[0]
        assert "dns.google" in bot.replied[0]


# -- Explicit type -----------------------------------------------------------

class TestExplicitType:
    def test_mx(self):
        bot = _FakeBot()
        mock = _mock_query(0, ["10 mail.example.com"])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com MX")))
        assert "MX:" in bot.replied[0]
        assert "mail.example.com" in bot.replied[0]

    def test_aaaa(self):
        bot = _FakeBot()
        mock = _mock_query(0, ["2606:2800:220:1:248:1893:25c8:1946"])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com AAAA")))
        assert "AAAA:" in bot.replied[0]

    def test_txt(self):
        bot = _FakeBot()
        mock = _mock_query(0, ["v=spf1 -all"])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com TXT")))
        assert "TXT:" in bot.replied[0]
        assert "v=spf1" in bot.replied[0]

    def test_case_insensitive(self):
        bot = _FakeBot()
        mock = _mock_query(0, ["93.184.216.34"])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com a")))
        assert "A:" in bot.replied[0]


# -- Unknown type ------------------------------------------------------------

class TestUnknownType:
    def test_unknown_type(self):
        bot = _FakeBot()
        asyncio.run(cmd_dns(bot, _msg("!dns example.com BOGUS")))
        assert "Unknown type" in bot.replied[0]


# -- Invalid PTR IP ----------------------------------------------------------

class TestInvalidPtr:
    def test_invalid_ip_for_ptr(self):
        bot = _FakeBot()
        asyncio.run(cmd_dns(bot, _msg("!dns notanip PTR")))
        assert "Invalid IP for PTR" in bot.replied[0]


# -- Query results -----------------------------------------------------------

class TestQueryResults:
    def test_success(self):
        bot = _FakeBot()
        mock = _mock_query(0, ["93.184.216.34", "93.184.216.35"])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com")))
        assert "93.184.216.34" in bot.replied[0]
        assert "93.184.216.35" in bot.replied[0]

    def test_timeout(self):
        bot = _FakeBot()
        mock = _mock_query(-1, [])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com")))
        assert "timeout" in bot.replied[0]

    def test_network_error(self):
        bot = _FakeBot()
        mock = _mock_query(-2, [])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com")))
        assert "network error" in bot.replied[0]

    def test_nxdomain(self):
        bot = _FakeBot()
        mock = _mock_query(3, [])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns nonexistent.test")))
        assert "NXDOMAIN" in bot.replied[0]

    def test_no_records(self):
        bot = _FakeBot()
        mock = _mock_query(0, [])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com AAAA")))
        assert "no records" in bot.replied[0]

    def test_servfail(self):
        bot = _FakeBot()
        mock = _mock_query(2, [])
        with patch.object(_mod, "_query", mock):
            asyncio.run(cmd_dns(bot, _msg("!dns example.com")))
        assert "SERVFAIL" in bot.replied[0]


# -- Missing args ------------------------------------------------------------

class TestMissingArgs:
    def test_no_args(self):
        bot = _FakeBot()
        asyncio.run(cmd_dns(bot, _msg("!dns")))
        assert "Usage" in bot.replied[0]