- Moved src/tests/ → tests/ - Added test_* functions with assertions to script-style test files - Added main() to each so IDEs offer it as a separate run target from pytest - Fixed cupy_test.py: remove spurious x_gpu += x_cpu, fix duplicate xlabel Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
91 lines
1.9 KiB
Python
91 lines
1.9 KiB
Python
import numpy as np
|
|
import pytest
|
|
from stack.rnn_helper import shift_right, shift_left, shift_up, shift_down, ch2idx, idx2ch, VOCAB
|
|
|
|
M = np.array([
|
|
[1, 2, 3],
|
|
[4, 5, 6],
|
|
], dtype=float)
|
|
|
|
|
|
def test_shift_right():
|
|
result = shift_right(M)
|
|
expected = np.array([
|
|
[0, 1, 2],
|
|
[0, 4, 5],
|
|
], dtype=float)
|
|
assert np.array_equal(result, expected)
|
|
assert result.shape == M.shape
|
|
|
|
|
|
def test_shift_left():
|
|
result = shift_left(M)
|
|
expected = np.array([
|
|
[2, 3, 0],
|
|
[5, 6, 0],
|
|
], dtype=float)
|
|
assert np.array_equal(result, expected)
|
|
assert result.shape == M.shape
|
|
|
|
|
|
def test_shift_up():
|
|
result = shift_up(M)
|
|
expected = np.array([
|
|
[4, 5, 6],
|
|
[0, 0, 0],
|
|
], dtype=float)
|
|
assert np.array_equal(result, expected)
|
|
assert result.shape == M.shape
|
|
|
|
|
|
def test_shift_down():
|
|
result = shift_down(M)
|
|
expected = np.array([
|
|
[0, 0, 0],
|
|
[1, 2, 3],
|
|
], dtype=float)
|
|
assert np.array_equal(result, expected)
|
|
assert result.shape == M.shape
|
|
|
|
|
|
def test_ch2idx_known():
|
|
assert ch2idx(' ') == 0
|
|
assert ch2idx('A') == 4
|
|
assert ch2idx('Z') == 29
|
|
assert ch2idx('0') == 30
|
|
assert ch2idx('9') == 39
|
|
|
|
|
|
def test_idx2ch_known():
|
|
assert idx2ch(0) == ' '
|
|
assert idx2ch(4) == 'A'
|
|
assert idx2ch(29) == 'Z'
|
|
assert idx2ch(30) == '0'
|
|
assert idx2ch(39) == '9'
|
|
|
|
|
|
def test_ch2idx_invalid():
|
|
with pytest.raises(KeyError):
|
|
ch2idx('a')
|
|
|
|
|
|
def test_idx2ch_invalid():
|
|
with pytest.raises(IndexError):
|
|
idx2ch(len(VOCAB))
|
|
|
|
|
|
def test_ch2idx_idx2ch_roundtrip():
|
|
for i, ch in enumerate(VOCAB):
|
|
assert ch2idx(ch) == i
|
|
assert idx2ch(i) == ch
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_shift_right()
|
|
test_shift_left()
|
|
test_shift_up()
|
|
test_shift_down()
|
|
test_ch2idx_known()
|
|
test_idx2ch_known()
|
|
test_ch2idx_idx2ch_roundtrip()
|
|
print("All tests passed.") |