def METHOD_NAME(self):
    with patch("salt.modules.dig.A", MagicMock(return_value=["ns4.google.com."])):
        dig_mock = MagicMock(
            return_value={
                "pid": 26136,
                "retcode": 0,
                "stderr": "",
                "stdout": "ns4.google.com.",
            }
        )
        with patch.dict(dig.__salt__, {"cmd.run_all": dig_mock}):
            self.assertEqual(dig.NS("google.com"), ["ns4.google.com."])