parent
eb280903f0
commit
123fc1e6c8
@ -0,0 +1,25 @@ |
||||
package selector |
||||
|
||||
import ( |
||||
"context" |
||||
) |
||||
|
||||
type peerKey struct{} |
||||
|
||||
// Peer contains the information of the peer for an RPC, such as the address
|
||||
// and authentication information.
|
||||
type Peer struct { |
||||
// node is the peer node.
|
||||
Node Node |
||||
} |
||||
|
||||
// NewPeerContext creates a new context with peer information attached.
|
||||
func NewPeerContext(ctx context.Context, p *Peer) context.Context { |
||||
return context.WithValue(ctx, peerKey{}, p) |
||||
} |
||||
|
||||
// FromPeerContext returns the peer information in ctx if it exists.
|
||||
func FromPeerContext(ctx context.Context) (p *Peer, ok bool) { |
||||
p, ok = ctx.Value(peerKey{}).(*Peer) |
||||
return |
||||
} |
@ -0,0 +1,24 @@ |
||||
package selector |
||||
|
||||
import ( |
||||
"context" |
||||
"testing" |
||||
) |
||||
|
||||
func TestPeer(t *testing.T) { |
||||
p := Peer{ |
||||
Node: mockWeightedNode{}, |
||||
} |
||||
ctx := NewPeerContext(context.Background(), &p) |
||||
p2, ok := FromPeerContext(ctx) |
||||
if !ok || p2.Node == nil { |
||||
t.Fatalf(" no peer found!") |
||||
} |
||||
} |
||||
|
||||
func TestNotPeer(t *testing.T) { |
||||
_, ok := FromPeerContext(context.Background()) |
||||
if ok { |
||||
t.Fatalf("test no peer found peer!") |
||||
} |
||||
} |
Loading…
Reference in new issue