import unittest
import numpy as np
from example import route, fixture
class RoutingTests(unittest.TestCase):
    def test_routes_and_normalization(self):
        x,r,e=fixture(); out,ids,w=route(x,r,e)
        np.testing.assert_array_equal(ids,[[0,1],[6,7]])
        np.testing.assert_allclose(w.sum(1),1)
        expected=w[0,0]*(x[0]@e[0])+w[0,1]*(x[0]@e[1])
        np.testing.assert_allclose(out[0],expected)
    def test_unused_expert_has_no_current_effect_but_still_stored(self):
        x,r,e=fixture(); before=route(x,r,e)[0]; original_size=e.nbytes
        e[3]+=1000
        np.testing.assert_array_equal(route(x,r,e)[0],before)
        self.assertEqual(e.nbytes,original_size)
        r[0,3]=100
        self.assertFalse(np.allclose(route(x,r,e)[0],before))
    def test_ties_are_stable(self):
        x,r,e=fixture(); r[:]=0
        np.testing.assert_array_equal(route(x,r,e)[1],[[0,1],[0,1]])
    def test_invalid_top_k(self):
        for k in (0,9,1.5,True):
            with self.assertRaises(ValueError): route(*fixture(),top_k=k)
    def test_bad_shapes_and_nonfinite(self):
        x,r,e=fixture()
        with self.assertRaises(ValueError): route(x,r[:2],e)
        x[0,0]=np.nan
        with self.assertRaises(ValueError): route(x,r,e)
if __name__=='__main__': unittest.main()