Add tests

pokedex/tests/test_struct.py

@@ -0,0 +1,162 @@
+# Encoding: utf8
+
+import base64
+
+import pytest
+
+from pokedex import struct
+from pokedex.db import connect, tables, util
+
+session = connect()
+
+def check_with_roundtrip(gen, pkmn, expected):
+    blob = pkmn.blob
+    del pkmn.blob
+    assert blob == pkmn.blob
+
+    assert pkmn.export_dict() == expected
+    from_dict = struct.save_file_pokemon_classes[5](session=session,
+        dict_=expected)
+    assert from_dict.blob == blob
+    assert from_dict.export_dict() == expected
+
+    from_blob = struct.save_file_pokemon_classes[5](session=session,
+        blob=pkmn.blob)
+    assert from_blob.blob == blob
+    assert from_blob.export_dict() == expected
+
+
+voltorb_species = util.get(session, tables.PokemonSpecies, 'voltorb')
+def voltorb_and_dict():
+    pkmn = struct.save_file_pokemon_classes[5](session=session)
+    voltorb_species = util.get(session, tables.PokemonSpecies, 'voltorb')
+    pkmn.species = voltorb_species
+    expected = {
+            'gender': 'male',
+            'species': dict(id=100, name=u'Voltorb'),
+            'nickname': u'\0' * 11,
+            'nickname trash': 'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==',
+            'moves': [],
+        }
+    return pkmn, expected
+
+
+def test_species():
+    pkmn, expected = voltorb_and_dict()
+    assert pkmn.species == voltorb_species
+    assert pkmn.pokemon == voltorb_species.default_pokemon
+    assert pkmn.form == voltorb_species.default_form
+    assert pkmn.export_dict() == expected
+
+
+def test_nickname():
+    pkmn, expected = voltorb_and_dict()
+    pkmn.nickname = pkmn.nickname
+    expected['nicknamed'] = True
+    check_with_roundtrip(5, pkmn, expected)
+
+    pkmn.is_nicknamed = False
+    del expected['nicknamed']
+    check_with_roundtrip(5, pkmn, expected)
+
+def test_moves():
+    pkmn, expected = voltorb_and_dict()
+    new_moves = (util.get(session, tables.Move, 'sonicboom'), )
+    expected['moves'] = [dict(id=49, name=u'SonicBoom', pp=0)]
+    pkmn.moves = new_moves
+    assert pkmn.moves == new_moves
+    check_with_roundtrip(5, pkmn, expected)
+
+    new_moves += (util.get(session, tables.Move, 'explosion'),)
+    expected['moves'].append(dict(id=153, name=u'Explosion', pp=0))
+    pkmn.moves = new_moves
+    assert pkmn.moves == new_moves
+    check_with_roundtrip(5, pkmn, expected)
+
+    new_pp = (20,)
+    expected['moves'][0]['pp'] = 20
+    pkmn.move_pp = new_pp
+    assert pkmn.move_pp == (20, 0, 0, 0)
+    check_with_roundtrip(5, pkmn, expected)
+
+def test_personality():
+    pkmn, expected = voltorb_and_dict()
+    assert pkmn.is_shiny == True
+    pkmn.personality = 12345
+    assert pkmn.is_shiny == False
+    expected['personality'] = 12345
+    check_with_roundtrip(5, pkmn, expected)
+
+def test_pokeball():
+    pkmn, expected = voltorb_and_dict()
+    masterball = util.get(session, tables.Item, 'master-ball')
+    pkmn.pokeball = masterball
+    assert pkmn.pokeball == masterball
+    expected['pokeball'] = dict(id_dppt=1, name='Master Ball')
+    check_with_roundtrip(5, pkmn, expected)
+
+def test_experience():
+    pkmn, expected = voltorb_and_dict()
+    exp = 2340
+    pkmn.exp = exp
+    assert pkmn.exp == exp
+    assert pkmn.experience_rung.experience < pkmn.exp
+    assert pkmn.next_experience_rung.experience > pkmn.exp
+    assert pkmn.experience_rung.level + 1 == pkmn.next_experience_rung.level
+    assert (pkmn.experience_rung.growth_rate ==
+        pkmn.next_experience_rung.growth_rate ==
+        pkmn.species.growth_rate)
+    assert pkmn.level == pkmn.experience_rung.level
+    assert pkmn.exp_to_next == pkmn.next_experience_rung.experience - pkmn.exp
+    rung_difference = (pkmn.next_experience_rung.experience -
+        pkmn.experience_rung.experience)
+    assert pkmn.progress_to_next == (
+        pkmn.exp - pkmn.experience_rung.experience) / float(rung_difference)
+    expected['exp'] = exp
+    check_with_roundtrip(5, pkmn, expected)
+
+def test_ability():
+    pkmn, expected = voltorb_and_dict()
+    ability = util.get(session, tables.Ability, 'drizzle')
+    pkmn.ability = ability
+    assert pkmn.ability == ability
+    expected['ability'] = dict(id=2, name='Drizzle')
+    check_with_roundtrip(5, pkmn, expected)
+
+def test_squitle_blob():
+    # Japanese Dream World Squirtle, coutresy of
+    # http://projectpokemon.org/events/
+    blob = base64.b64decode('J2ZqBgAAICQHAAAAkOaKyTACAABGLAABAAAAAAAAAAAAAAAAA'
+        'AAAACEAJwCRAG4AIx4eKAAAAAD171MHAAAAAAAAAQAAAAAAvDDLMKww4TD//wAAAAAAAA'
+        'AAAAD//wAVAAAAAAAAAAAw/zD/T/9S/0f///8AAAAAAAAACgoOAABLAAAZCgAAAA==')
+    expected = {
+        'ability': {'id': 44, 'name': u'Rain Dish'},
+        'date met': '2010-10-14',
+        'exp': 560,
+        'gender': 'male',
+        'genes': {u'attack': 31,
+                u'defense': 27,
+                u'hp': 21,
+                u'special attack': 21,
+                u'special defense': 3,
+                u'speed': 7},
+        'happiness': 70,
+        'met at level': 10,
+        'met location': {'id_dp': 75, 'name': u'Spring Path'},
+        'moves': [{'id': 33, 'name': u'Tackle', 'pp': 35},
+                {'id': 39, 'name': u'Tail Whip', 'pp': 30},
+                {'id': 145, 'name': u'Bubble', 'pp': 30},
+                {'id': 110, 'name': u'Withdraw', 'pp': 40}],
+        'nickname': u'ゼニガメ',
+        'nickname trash': 'vDDLMKww4TD//wAAAAAAAAAAAAD//w==',
+        'oiginal trainer': {'gender': 'male',
+                            'id': 59024,
+                            'name': u'ＰＰｏｒｇ',
+                            'secret': 51594},
+        'original country': 'jp',
+        'original version': 21,
+        'personality': 107636263,
+        'pokeball': {'id_dppt': 25, 'name': u'Hyper Potion'},
+        'species': {'id': 7, 'name': u'Squirtle'}}
+    pkmn = struct.save_file_pokemon_classes[5](session=session, blob=blob)
+    check_with_roundtrip(5, pkmn, expected)

scripts/test-struct-roundtrip.py

@@ -0,0 +1,125 @@
+#! /usr/bin/env python
+"""
+This is an ad-hoc testing script. YMMV
+"""
+
+import os
+import sys
+import pprint
+import binascii
+import traceback
+import subprocess
+import tempfile
+import itertools
+
+import yaml  # you need to pip install pyyaml
+from blessings import Terminal  # you need to pip install blessings
+
+from pokedex import struct
+from pokedex.db import connect
+
+session = connect(engine_args=dict(echo=False))
+
+if len(sys.argv) < 1:
+    print 'Give this script a bunch of PKM files to test roundtrips on.'
+    print 'A number (e.g. "4") will be interpreted as the generation of'
+    print 'the following files, until a new generation is given.'
+    print 'Use "./5" for a file named 5.'
+    print
+    print 'If mismatches are found, your screen will be filled with colorful'
+    print 'reports. You need the colordiff program for this.'
+
+def printable(c):
+    if ord(' ') < ord(c) < ord('~'):
+        return c
+    else:
+        return '.'
+
+def colordiff(str1, str2, prefix='tmp-'):
+    if str1 != str2:
+        with tempfile.NamedTemporaryFile(prefix=prefix + '.', suffix='.a') as file1:
+            with tempfile.NamedTemporaryFile(prefix=prefix + '.', suffix='.b') as file2:
+                file1.write(str1)
+                file2.write(str2)
+                file1.flush()
+                file2.flush()
+                p = subprocess.Popen(['colordiff', '-U999', file1.name, file2.name])
+                p.communicate()
+    else:
+        print prefix, 'match:'
+        print str1
+
+Class = struct.save_file_pokemon_classes[5]
+
+filenames_left = list(reversed(sys.argv[1:]))
+
+while filenames_left:
+    filename = filenames_left.pop()
+    print filename
+
+    try:
+        generation = int(filename)
+    except ValueError:
+        pass
+    else:
+        Class = struct.save_file_pokemon_classes[generation]
+        continue
+
+    if os.path.isdir(filename):
+        for name in sorted(os.listdir(filename), reverse=True):
+            joined = os.path.join(filename, name)
+            if name.endswith('.pkm') or os.path.isdir(joined):
+                filenames_left.append(joined)
+        continue
+
+    with open(filename) as f:
+        blob = f.read()[:0x88]
+
+    if blob[0] == blob[1] == blob[2] == blob[3] == '\0':
+        print binascii.hexlify(blob)
+        print 'Probably not a PKM file'
+
+    try:
+        orig_object = Class(blob, session=session)
+        dict_ = orig_object.export_dict()
+    except Exception:
+        traceback.print_exc()
+        print binascii.hexlify(blob)
+        continue
+    orig_object.blob
+    new_object = Class(dict_=dict_, session=session)
+    try:
+        blob_again = new_object.blob
+        dict_again = new_object.export_dict()
+    except Exception:
+        colordiff(yaml.dump(orig_object.structure), yaml.dump(new_object.structure), 'struct')
+        traceback.print_exc()
+        continue
+
+    if (dict_ != dict_again) or (blob != blob_again):
+        colordiff(yaml.dump(orig_object.structure), yaml.dump(new_object.structure), 'struct')
+        colordiff(yaml.safe_dump(dict_), yaml.safe_dump(dict_again), 'yaml')
+        t = Terminal()
+        for pass_number in 1, 2, 3:
+            for i, (a, b) in enumerate(itertools.izip_longest(blob, blob_again, fillvalue='\xbb')):
+                if (i - 8) % 32 == 0:
+                    # Block boundary
+                    sys.stdout.write(' ')
+                a_hex = binascii.hexlify(a)
+                b_hex = binascii.hexlify(b)
+                if a != b:
+                    if pass_number == 1:
+                        sys.stdout.write(t.green(printable(a)))
+                        sys.stdout.write(t.red(printable(b)))
+                    elif pass_number == 2:
+                        sys.stdout.write(t.green(a_hex))
+                    elif pass_number == 3:
+                        sys.stdout.write(t.red(b_hex))
+                else:
+                    if pass_number == 1:
+                        sys.stdout.write(printable(a))
+                        sys.stdout.write(printable(b))
+                    else:
+                        sys.stdout.write(a_hex)
+            print
+        print