github.com/whoyao/protocol@v0.0.0-20230519045905-2d8ace718ca5/utils/graph_test.go (about)

     1  package utils
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/require"
     7  	"golang.org/x/exp/maps"
     8  )
     9  
    10  type testNode struct {
    11  	id string
    12  }
    13  
    14  func (n *testNode) ID() string {
    15  	return n.id
    16  }
    17  
    18  type testEdge int
    19  
    20  func (e testEdge) Length() int64 {
    21  	return int64(e)
    22  }
    23  
    24  func testNodeIDs[K comparable, N GraphNodeProps[K]](path []N) []K {
    25  	var names []K
    26  	for _, p := range path {
    27  		names = append(names, p.ID())
    28  	}
    29  	return names
    30  }
    31  
    32  func TestGraph(t *testing.T) {
    33  	t.Run("graph mutation", func(t *testing.T) {
    34  		g := NewGraph[string, *testNode, testEdge]()
    35  		g.InsertNode(&testNode{"a"})
    36  		g.InsertNode(&testNode{"b"})
    37  		g.InsertNode(&testNode{"c"})
    38  
    39  		require.EqualValues(t, &testNode{"a"}, g.Node("a"))
    40  
    41  		g.InsertEdge("a", "b", 1)
    42  		g.InsertEdge("b", "c", 1)
    43  		g.InsertEdge("c", "a", 1)
    44  
    45  		require.Equal(t, []string{"b"}, maps.Keys(g.OutEdges("a")))
    46  		require.Equal(t, []string{"c"}, maps.Keys(g.InEdges("a")))
    47  
    48  		e, ok := g.Edge("a", "b")
    49  		require.EqualValues(t, 1, e)
    50  		require.True(t, ok)
    51  
    52  		g.DeleteEdge("a", "b")
    53  		_, ok = g.Edge("a", "b")
    54  		require.False(t, ok)
    55  	})
    56  
    57  	t.Run("topological sort", func(t *testing.T) {
    58  		g := NewGraph[string, *testNode, testEdge]()
    59  		g.InsertNode(&testNode{"a"})
    60  		g.InsertNode(&testNode{"b"})
    61  		g.InsertNode(&testNode{"c"})
    62  		g.InsertNode(&testNode{"d"})
    63  		g.InsertNode(&testNode{"e"})
    64  
    65  		g.InsertEdge("a", "b", 1)
    66  		g.InsertEdge("b", "c", 1)
    67  		g.InsertEdge("c", "d", 1)
    68  		g.InsertEdge("d", "e", 1)
    69  
    70  		require.Equal(t, []string{"a", "b", "c", "d", "e"}, testNodeIDs[string](g.TopologicalSort()))
    71  	})
    72  
    73  	t.Run("shortest path", func(t *testing.T) {
    74  		g := NewGraph[string, *testNode, testEdge]()
    75  		g.InsertNode(&testNode{"a"})
    76  		g.InsertNode(&testNode{"b"})
    77  		g.InsertNode(&testNode{"c"})
    78  		g.InsertNode(&testNode{"d"})
    79  
    80  		g.InsertEdge("a", "b", 1)
    81  		g.InsertEdge("a", "c", 3)
    82  		g.InsertEdge("b", "d", 2)
    83  		g.InsertEdge("c", "d", 1)
    84  
    85  		p, n := g.ShortestPath("a", "d")
    86  		require.Equal(t, []string{"a", "b", "d"}, testNodeIDs[string](p))
    87  		require.EqualValues(t, 3, n)
    88  	})
    89  }