demand-agent-new/examples/demand/tests/test_data_query_tools.py

import unittest

from examples.demand.data_query_tools import _requested_count_from_lack_count


class DataQueryToolsTest(unittest.TestCase):
    def test_requested_count_from_lack_count_uses_reduced_buckets(self):
        cases = [
            (1000.01, 7),
            (1000, 6),
            (500.01, 6),
            (500, 4),
            (100.01, 4),
            (100, 2),
            (50.01, 2),
            (50, 1),
            (0, 1),
        ]

        for lack_count, expected in cases:
            with self.subTest(lack_count=lack_count):
                self.assertEqual(_requested_count_from_lack_count(lack_count), expected)


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