From daeb75bac16cdcc9d6108e889740c33c52e7d2cb Mon Sep 17 00:00:00 2001
From: Tao <tao.li06@sap.com>
Date: Thu, 23 Nov 2023 15:23:39 +0100
Subject: [PATCH] add function to get subscribed vnis as a client

---
 metalbond.go | 15 +++++++++++++--
 peer_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/metalbond.go b/metalbond.go
index b9490ab..4253ddb 100644
--- a/metalbond.go
+++ b/metalbond.go
@@ -40,8 +40,7 @@ type MetalBond struct {
 
 	keepaliveInterval uint32
 	shuttingDown      bool
-
-	client Client
+	client            Client
 
 	lis      *net.Listener // for server only
 	isServer bool
@@ -310,6 +309,18 @@ func (m *MetalBond) GetRoutesForVni(vni VNI) error {
 	return nil
 }
 
+func (m *MetalBond) GetSubscribedVnis() []VNI {
+	m.mtxMySubscriptions.RLock()
+	defer m.mtxMySubscriptions.RUnlock()
+
+	vnis := make([]VNI, 0)
+	for vni := range m.mySubscriptions {
+		vnis = append(vnis, vni)
+	}
+
+	return vnis
+}
+
 func (m *MetalBond) addReceivedRoute(fromPeer *metalBondPeer, vni VNI, dest Destination, hop NextHop) error {
 	err := m.routeTable.AddNextHop(vni, dest, hop, fromPeer)
 	if err != nil {
diff --git a/peer_test.go b/peer_test.go
index 33c6c12..4ddb65b 100644
--- a/peer_test.go
+++ b/peer_test.go
@@ -144,6 +144,42 @@ var _ = Describe("Peer", func() {
 		Expect(err).NotTo(HaveOccurred())
 	})
 
+	It("should subscribe", func() {
+		mbClient := NewMetalBond(Config{}, client)
+		localIP := net.ParseIP("127.0.0.1")
+		err := mbClient.AddPeer(serverAddress, localIP.String())
+		Expect(err).NotTo(HaveOccurred())
+
+		time.Sleep(1 * time.Second)
+		vni_one := VNI(200)
+		err = mbClient.Subscribe(vni_one)
+		if err != nil {
+			log.Errorf("subscribe failed: %v", err)
+		}
+		Expect(err).NotTo(HaveOccurred())
+
+		vni_two := VNI(300)
+		err = mbClient.Subscribe(vni_two)
+		if err != nil {
+			log.Errorf("subscribe failed: %v", err)
+		}
+		Expect(err).NotTo(HaveOccurred())
+
+		vnis := mbClient.GetSubscribedVnis()
+		Expect(len(vnis)).To(Equal(2))
+		Expect(vnis[0]).To(Equal(vni_one))
+		Expect(vnis[1]).To(Equal(vni_two))
+
+		err = mbClient.Unsubscribe(vni_one)
+		Expect(err).NotTo(HaveOccurred())
+
+		err = mbClient.Unsubscribe(vni_two)
+		Expect(err).NotTo(HaveOccurred())
+
+		vnis = mbClient.GetSubscribedVnis()
+		Expect(len(vnis)).To(Equal(0))
+	})
+
 	It("should announce", func() {
 		totalClients := 1000
 		var wg sync.WaitGroup