Spaces:

sitcod3
/

modul-x

Running

modul-x / tests /test_engine_ml.py

Upload folder using huggingface_hub

d82b7e3 verified 8 days ago

8.04 kB

import unittest
from unittest.mock import AsyncMock, MagicMock, patch

from app.models import DNSResult, DisposableResult, SMTPResult, ValidationStatus
from app.validator.engine import EmailValidationEngine


def _dns():
    return DNSResult(
        domain_exists=True,
        has_mx_records=True,
        mx_records=[{"host": "mx.corp.test", "priority": 10}],
    )


def _clean():
    return DisposableResult(
        is_disposable=False, is_free_provider=False, is_role_based=False
    )


class NullCache:
    async def get_dns(self, domain):
        return None

    async def put_dns(self, domain, result):
        return None

    async def get_verdict(self, email, deep):
        return None

    async def put_verdict(self, email, deep, smtp, catch_all):
        return None


class FakeProvider:
    async def verify_many(self, pairs):
        return ({e: SMTPResult(verified=True, smtp_code=250) for e, _ in pairs}, {})

    async def verify(self, email, mx_host):
        return SMTPResult(verified=True, smtp_code=250)

    async def catch_all(self, domain, mx_host):
        return (False, 0.9)


def _engine_with_ml(flag=True, repo="user/model", predict=None):
    engine = EmailValidationEngine(mailbox=FakeProvider(), cache=NullCache())
    # The get_model patcher must honor the flag: with the flag off the REAL
    # get_model returns None (that IS the production contract), so the double
    # must not force a model into a disabled engine.
    if predict is not None and flag:
        fake_model = MagicMock()
        fake_model.predict = MagicMock(side_effect=predict)
        fake_model.load = MagicMock(return_value=True)
        engine.mailbox = engine.mailbox  # unchanged
        patcher = patch("app.validator.engine.get_model_async",
                        new=AsyncMock(return_value=fake_model))
        patcher.start()
        engine._ml_patcher = patcher  # keep alive until stopAll
        engine._ml_model = fake_model
    p = patch.object(
        __import__("app.config", fromlist=["settings"]).settings,
        "ml_disposable_enabled",
        flag,
    )
    p.start()
    engine._flag_patcher = p
    p2 = patch.object(
        __import__("app.config", fromlist=["settings"]).settings,
        "ml_model_repo",
        repo,
    )
    p2.start()
    engine._flag_patcher2 = p2
    return engine


def _stop(engine):
    engine._ml_patcher.stop() if hasattr(engine, "_ml_patcher") else None
    engine._flag_patcher.stop()
    engine._flag_patcher2.stop()


class EngineMLTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        self._engines = []

    def tearDown(self):
        # Stop settings/model patchers so they never leak into other tests.
        for e in self._engines:
            _stop(e)
    def tearDown(self):
        for e in getattr(self, "_engines", []):
            _stop(e)

    async def _run(self, engine, emails):
        self._engines = getattr(self, "_engines", [])
        self._engines.append(engine)
        # NOTE: check_disposable is NOT mocked - the real blocklist/role sets
        # run, so the known-domain test exercises the true deterministic path.
        with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
            return await engine.validate_batch(emails, deep=True)

    async def test_flag_off_leaves_results_unchanged(self):
        engine = _engine_with_ml(flag=False)
        results = await self._run(engine, ["a@brand-new-xyz.tk"])

        self.assertIsNone(results[0].disposable.ml_label)
        self.assertIsNone(results[0].disposable.ml_confidence)
        self.assertNotIn(
            "ML model flagged domain as disposable-like",
            results[0].scoring.warnings,
        )

    async def test_known_disposable_domain_skips_ml(self):
        # mailinator.com IS in the live blocklist -> the deterministic layer
        # already hard-fails it; the model must not be consulted.
        engine = _engine_with_ml(flag=True, predict=[AssertionError("must not run")])
        results = await self._run(engine, ["a@mailinator.com"])

        self.assertEqual(results[0].status, ValidationStatus.INVALID)
        fake_model = engine._ml_model
        fake_model.predict.assert_not_called()

    async def test_unknown_domain_gets_advisory_ml_signal(self):
        engine = _engine_with_ml(
            flag=True, predict=[[("DISPOSABLE", 0.93)]]  # contract: (label, score)
        )
        results = await self._run(engine, ["a@brand-new-disposable-xyz.tk"])

        d = results[0].disposable
        self.assertEqual(d.ml_label, "DISPOSABLE")
        self.assertEqual(d.ml_confidence, 0.93)
        self.assertTrue(
            any("ML model flagged" in w for w in results[0].scoring.warnings)
        )
        # ADVISORY: the deterministic verdict must not become invalid.
        self.assertNotEqual(results[0].status, ValidationStatus.INVALID)
        # The penalty is proven by the breakdown: the SMTP-verified bonus (+10)
        # can legitimately re-inflate the absolute score past the clamp.
        self.assertEqual(
            results[0].scoring.breakdown.get("ML_DISPOSABLE_SUSPECT"), -5
        )

    async def test_model_failure_degrades_to_deterministic(self):
        engine = _engine_with_ml(flag=True, predict=[None])
        results = await self._run(engine, ["a@brand-new-xyz.tk"])

        self.assertIsNone(results[0].disposable.ml_label)
        self.assertEqual(
            results[0].scoring.warnings,
            [w for w in results[0].scoring.warnings if "ML model flagged" not in w]
            or results[0].scoring.warnings,
        )
        self.assertEqual(results[0].scoring.score, 100)

    async def test_unknown_domains_batch_into_one_predict_call(self):
        engine = _engine_with_ml(
            flag=True,
            predict=[
                [
                    ("DISPOSABLE", 0.9),
                    ("NORMAL", 0.9),
                    ("DISPOSABLE", 0.95),
                ]
            ],
        )
        emails = [f"u{i}@dk{i}.example" for i in range(3)]
        results = await self._run(engine, emails)

        fake_model = engine._ml_model
        self.assertEqual(fake_model.predict.call_count, 1)
        sent = fake_model.predict.call_args.args[0]
        self.assertEqual(len(sent), 3)
        self.assertEqual(results[1].disposable.ml_label, "NORMAL")


    async def test_single_validate_also_applies_the_ml_signal(self):
        # The single endpoint runs engine.validate() - it must get the same
        # advisory signal as bulk, or the two surfaces disagree on one address.
        engine = _engine_with_ml(
            flag=True, predict=[[("DISPOSABLE", 0.93)]]
        )
        self._engines.append(engine)
        with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
            result = await engine.validate("a@brand-new-disposable-xyz.tk", deep=True)

        self.assertEqual(result.disposable.ml_label, "DISPOSABLE")
        self.assertEqual(
            result.scoring.breakdown.get("ML_DISPOSABLE_SUSPECT"), -5
        )

    async def test_single_validate_flag_off_unchanged(self):
        engine = _engine_with_ml(flag=False, predict=[[("DISPOSABLE", 0.93)]])
        self._engines.append(engine)
        with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
            result = await engine.validate("a@brand-new-disposable-xyz.tk", deep=True)

        self.assertIsNone(result.disposable.ml_label)

    async def test_single_validate_known_domain_skips_ml(self):
        engine = _engine_with_ml(
            flag=True, predict=[AssertionError("must not run")]
        )
        self._engines.append(engine)
        with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
            result = await engine.validate("a@mailinator.com", deep=True)

        self.assertEqual(result.status, ValidationStatus.INVALID)
        self.assertIsNone(result.disposable.ml_label)


if __name__ == "__main__":
    unittest.main()

	import unittest
	from unittest.mock import AsyncMock, MagicMock, patch

	from app.models import DNSResult, DisposableResult, SMTPResult, ValidationStatus
	from app.validator.engine import EmailValidationEngine


	def _dns():
	return DNSResult(
	domain_exists=True,
	has_mx_records=True,
	mx_records=[{"host": "mx.corp.test", "priority": 10}],
	)


	def _clean():
	return DisposableResult(
	is_disposable=False, is_free_provider=False, is_role_based=False
	)


	class NullCache:
	async def get_dns(self, domain):
	return None

	async def put_dns(self, domain, result):
	return None

	async def get_verdict(self, email, deep):
	return None

	async def put_verdict(self, email, deep, smtp, catch_all):
	return None


	class FakeProvider:
	async def verify_many(self, pairs):
	return ({e: SMTPResult(verified=True, smtp_code=250) for e, _ in pairs}, {})

	async def verify(self, email, mx_host):
	return SMTPResult(verified=True, smtp_code=250)

	async def catch_all(self, domain, mx_host):
	return (False, 0.9)


	def _engine_with_ml(flag=True, repo="user/model", predict=None):
	engine = EmailValidationEngine(mailbox=FakeProvider(), cache=NullCache())
	# The get_model patcher must honor the flag: with the flag off the REAL
	# get_model returns None (that IS the production contract), so the double
	# must not force a model into a disabled engine.
	if predict is not None and flag:
	fake_model = MagicMock()
	fake_model.predict = MagicMock(side_effect=predict)
	fake_model.load = MagicMock(return_value=True)
	engine.mailbox = engine.mailbox # unchanged
	patcher = patch("app.validator.engine.get_model_async",
	new=AsyncMock(return_value=fake_model))
	patcher.start()
	engine._ml_patcher = patcher # keep alive until stopAll
	engine._ml_model = fake_model
	p = patch.object(
	__import__("app.config", fromlist=["settings"]).settings,
	"ml_disposable_enabled",
	flag,
	)
	p.start()
	engine._flag_patcher = p
	p2 = patch.object(
	__import__("app.config", fromlist=["settings"]).settings,
	"ml_model_repo",
	repo,
	)
	p2.start()
	engine._flag_patcher2 = p2
	return engine


	def _stop(engine):
	engine._ml_patcher.stop() if hasattr(engine, "_ml_patcher") else None
	engine._flag_patcher.stop()
	engine._flag_patcher2.stop()


	class EngineMLTests(unittest.IsolatedAsyncioTestCase):
	def setUp(self):
	self._engines = []

	def tearDown(self):
	# Stop settings/model patchers so they never leak into other tests.
	for e in self._engines:
	_stop(e)
	def tearDown(self):
	for e in getattr(self, "_engines", []):
	_stop(e)

	async def _run(self, engine, emails):
	self._engines = getattr(self, "_engines", [])
	self._engines.append(engine)
	# NOTE: check_disposable is NOT mocked - the real blocklist/role sets
	# run, so the known-domain test exercises the true deterministic path.
	with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
	return await engine.validate_batch(emails, deep=True)

	async def test_flag_off_leaves_results_unchanged(self):
	engine = _engine_with_ml(flag=False)
	results = await self._run(engine, ["a@brand-new-xyz.tk"])

	self.assertIsNone(results[0].disposable.ml_label)
	self.assertIsNone(results[0].disposable.ml_confidence)
	self.assertNotIn(
	"ML model flagged domain as disposable-like",
	results[0].scoring.warnings,
	)

	async def test_known_disposable_domain_skips_ml(self):
	# mailinator.com IS in the live blocklist -> the deterministic layer
	# already hard-fails it; the model must not be consulted.
	engine = _engine_with_ml(flag=True, predict=[AssertionError("must not run")])
	results = await self._run(engine, ["a@mailinator.com"])

	self.assertEqual(results[0].status, ValidationStatus.INVALID)
	fake_model = engine._ml_model
	fake_model.predict.assert_not_called()

	async def test_unknown_domain_gets_advisory_ml_signal(self):
	engine = _engine_with_ml(
	flag=True, predict=[[("DISPOSABLE", 0.93)]] # contract: (label, score)
	)
	results = await self._run(engine, ["a@brand-new-disposable-xyz.tk"])

	d = results[0].disposable
	self.assertEqual(d.ml_label, "DISPOSABLE")
	self.assertEqual(d.ml_confidence, 0.93)
	self.assertTrue(
	any("ML model flagged" in w for w in results[0].scoring.warnings)
	)
	# ADVISORY: the deterministic verdict must not become invalid.
	self.assertNotEqual(results[0].status, ValidationStatus.INVALID)
	# The penalty is proven by the breakdown: the SMTP-verified bonus (+10)
	# can legitimately re-inflate the absolute score past the clamp.
	self.assertEqual(
	results[0].scoring.breakdown.get("ML_DISPOSABLE_SUSPECT"), -5
	)

	async def test_model_failure_degrades_to_deterministic(self):
	engine = _engine_with_ml(flag=True, predict=[None])
	results = await self._run(engine, ["a@brand-new-xyz.tk"])

	self.assertIsNone(results[0].disposable.ml_label)
	self.assertEqual(
	results[0].scoring.warnings,
	[w for w in results[0].scoring.warnings if "ML model flagged" not in w]
	or results[0].scoring.warnings,
	)
	self.assertEqual(results[0].scoring.score, 100)

	async def test_unknown_domains_batch_into_one_predict_call(self):
	engine = _engine_with_ml(
	flag=True,
	predict=[
	[
	("DISPOSABLE", 0.9),
	("NORMAL", 0.9),
	("DISPOSABLE", 0.95),
	]
	],
	)
	emails = [f"u{i}@dk{i}.example" for i in range(3)]
	results = await self._run(engine, emails)

	fake_model = engine._ml_model
	self.assertEqual(fake_model.predict.call_count, 1)
	sent = fake_model.predict.call_args.args[0]
	self.assertEqual(len(sent), 3)
	self.assertEqual(results[1].disposable.ml_label, "NORMAL")


	async def test_single_validate_also_applies_the_ml_signal(self):
	# The single endpoint runs engine.validate() - it must get the same
	# advisory signal as bulk, or the two surfaces disagree on one address.
	engine = _engine_with_ml(
	flag=True, predict=[[("DISPOSABLE", 0.93)]]
	)
	self._engines.append(engine)
	with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
	result = await engine.validate("a@brand-new-disposable-xyz.tk", deep=True)

	self.assertEqual(result.disposable.ml_label, "DISPOSABLE")
	self.assertEqual(
	result.scoring.breakdown.get("ML_DISPOSABLE_SUSPECT"), -5
	)

	async def test_single_validate_flag_off_unchanged(self):
	engine = _engine_with_ml(flag=False, predict=[[("DISPOSABLE", 0.93)]])
	self._engines.append(engine)
	with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
	result = await engine.validate("a@brand-new-disposable-xyz.tk", deep=True)

	self.assertIsNone(result.disposable.ml_label)

	async def test_single_validate_known_domain_skips_ml(self):
	engine = _engine_with_ml(
	flag=True, predict=[AssertionError("must not run")]
	)
	self._engines.append(engine)
	with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
	result = await engine.validate("a@mailinator.com", deep=True)

	self.assertEqual(result.status, ValidationStatus.INVALID)
	self.assertIsNone(result.disposable.ml_label)


	if __name__ == "__main__":
	unittest.main()