<commit_before>
from socket import inet_aton
import tempfile

from whip.db import Database
from whip.json import loads as json_loads
from whip.util import ipv4_str_to_int


def test_db_loading():

    def t(begin, end, x, datetime):
        """Helper to create test data"""
        return dict(begin=begin, end=end, x=x, datetime=datetime)

    snapshots = [
        [
            # Initial data
            t('1.0.0.0', '172.16.31.10', 1, '2010'),
            t('3.0.0.0', '192.168.3.11', 2, '2010'),
            t('8.0.0.0', '192.168.3.11', 3, '2010'),
        ],
        [
            # Split some ranges, exclude some ranges
            t('1.0.0.0', '1.2.3.4', 7, '2011'),
            t('1.2.3.5', '1.3.4.5', 8, '2011'),
        ],
        [
            # Merge some ranges, update some values
            t('1.0.0.0', '172.16.31.10', 4, '2013'),
            t('3.0.0.0', '192.168.3.11', 5, '2013'),
            t('8.0.0.0', '192.168.3.11', 6, '2013'),
        ],
    ]

    def iter_snapshot(snapshot):
        for d in snapshot:
            yield ipv4_str_to_int(d['begin']), ipv4_str_to_int(d['end']), d

    with tempfile.TemporaryDirectory() as db_dir:
        db = Database(db_dir, create_if_missing=True)
        iters = [iter_snapshot(s) for s in snapshots]

        # db.load(*iters)

        db.load(*iters[:-1])
        db.load(iters[-1])
        db.load()

        def lookup(ip, datetime=None):
            """Lookup a single version"""
            ret = db.lookup(inet_aton(ip), datetime=datetime) or b'{}'
            return json_loads(ret)

        def lookup_all_x(ip):
            """Lookup all versions, returning only the 'x' values"""
            history = lookup(ip, 'all')['history']
            return [d['x'] for d in history]

        # Latest version
        assert lookup('1.0.0.0')['x'] == 4
        assert lookup('172.16.31.10')['x'] == 4
        assert lookup('7.0.0.0') == {}
        assert lookup('8.1.2.3')['x'] == 6
        assert lookup('1172.16.17.32') == {}

        # Specific dates
        assert lookup('1.2.3.3', '2010')['x'] == 1
        assert lookup('1.2.3.4', '2011')['x'] == 7
        assert lookup('1.2.3.5', '2011')['x'] == 8
        assert lookup('1.100.100.100', '2011')['x'] == 1
        assert lookup('8.1.2.3', '2011')['x'] == 3

        # No hit for really old dates
        assert lookup('1.2.3.4', '2009') == {}

        # Future date
        assert lookup('1.2.3.4', '2038')['x'] == 4

        # All versions
        assert lookup_all_x('1.2.3.4') == [4, 7, 1]<commit_msg>Some cleanups in unit tests<commit_after>
from socket import inet_aton
import tempfile

from whip.db import Database
from whip.json import loads as json_loads
from whip.util import ipv4_str_to_int


def test_db_loading():

    def t(begin, end, x, datetime):
        """Helper to create test data"""
        return dict(begin=begin, end=end, x=x, datetime=datetime)

    def iter_snapshot(snapshot):
        for d in snapshot:
            yield ipv4_str_to_int(d['begin']), ipv4_str_to_int(d['end']), d

    def lookup(db, ip, datetime=None):
        """Lookup a single version"""
        ret = db.lookup(inet_aton(ip), datetime=datetime) or b'{}'
        return json_loads(ret)

    def lookup_all_x(db, ip):
        """Lookup all versions, returning only the 'x' values"""
        history = lookup(db, ip, 'all')['history']
        return [d['x'] for d in history]

    snapshots = [
        [
            # Initial data
            t('1.0.0.0', '172.16.31.10', 1, '2010'),
            t('3.0.0.0', '192.168.3.11', 2, '2010'),
            t('8.0.0.0', '192.168.3.11', 3, '2010'),
        ],
        [
            # Split some ranges, exclude some ranges
            t('1.0.0.0', '1.2.3.4', 7, '2011'),
            t('1.2.3.5', '1.3.4.5', 8, '2011'),
        ],
        [
            # Merge some ranges, update some values
            t('1.0.0.0', '172.16.31.10', 4, '2013'),
            t('3.0.0.0', '192.168.3.11', 5, '2013'),
            t('8.0.0.0', '192.168.3.11', 6, '2013'),
        ],
    ]

    def check_sanity(db):

        # Latest version
        assert lookup(db, '1.0.0.0')['x'] == 4
        assert lookup(db, '172.16.31.10')['x'] == 4
        assert lookup(db, '7.0.0.0') == {}
        assert lookup(db, '8.1.2.3')['x'] == 6
        assert lookup(db, '1172.16.17.32') == {}

        # Specific dates
        assert lookup(db, '1.2.3.3', '2010')['x'] == 1
        assert lookup(db, '1.2.3.4', '2011')['x'] == 7
        assert lookup(db, '1.2.3.5', '2011')['x'] == 8
        assert lookup(db, '172.16.31.10', '2011')['x'] == 1
        assert lookup(db, '8.1.2.3', '2011')['x'] == 3

        # No hit for really old dates
        assert lookup(db, '1.2.3.4', '2009') == {}

        # Future date
        assert lookup(db, '1.2.3.4', '2038')['x'] == 4

        # All versions
        assert lookup_all_x(db, '1.2.3.4') == [4, 7, 1]

    with tempfile.TemporaryDirectory() as db_dir:
        db = Database(db_dir, create_if_missing=True)
        iters = [iter_snapshot(s) for s in snapshots]

        # db.load(*iters)

        db.load(*iters[:-1])
        db.load(iters[-1])
        db.load(iters[-1])
        db.load()

        check_sanity(db)
{'fraction_non_alphanumeric': 0.19275948392066242, 'fraction_numerical': 0.09551319083381475, 'mean_word_length': 1.9528140989198408, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 58, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6390221', 'n_tokens_mistral': 2127, 'n_tokens_neox': 1881, 'n_words': 431}