From 6eab4f0b297fd98c6a416987ef482fda2887ef36 Mon Sep 17 00:00:00 2001
From: Ioannis Simeonidis <36934072+simioa@users.noreply.github.com>
Date: Wed, 18 Dec 2024 09:06:31 +0100
Subject: [PATCH] add a test that tests the template sharing better

---
 .../input/netflow/decoder/v9/v9_test.go       | 50 ++++++++-----------
 1 file changed, 22 insertions(+), 28 deletions(-)

diff --git a/x-pack/filebeat/input/netflow/decoder/v9/v9_test.go b/x-pack/filebeat/input/netflow/decoder/v9/v9_test.go
index 9134528c673e..67212c1e4084 100644
--- a/x-pack/filebeat/input/netflow/decoder/v9/v9_test.go
+++ b/x-pack/filebeat/input/netflow/decoder/v9/v9_test.go
@@ -251,8 +251,8 @@ func TestCustomFields(t *testing.T) {
 }
 
 func TestSharedTemplates(t *testing.T) {
-	const sourceID = 1234
-	addr := test.MakeAddress(t, "127.0.0.1:12345")
+	templateAddr := test.MakeAddress(t, "127.0.0.1:12345")
+	flowsAddr := test.MakeAddress(t, "127.0.0.2:21234")
 	templatePacket := []uint16{
 		// Header
 		// Version, Count, Uptime, Ts, SeqNo, Source
@@ -264,44 +264,38 @@ func TestSharedTemplates(t *testing.T) {
 		2, 4,
 		3, 4,
 	}
+	flowsPacket := []uint16{
+		// Header
+		// Version, Count, Uptime, Ts, SeqNo, Source
+		9, 1, 11, 11, 22, 22, 33, 34, 0, 1234,
+		// Set #1 (template)
+		999, 16, /*len of set*/
+		1, 1,
+		2, 2,
+		3, 3,
+	}
 
 	t.Run("Template sharing enabled", func(t *testing.T) {
-		sharedTemplates := true
-		key := MakeSessionKey(addr, sourceID, sharedTemplates)
 		cfg := config.Defaults()
-		cfg.WithSharedTemplates(sharedTemplates)
+		cfg.WithSharedTemplates(true)
 		proto := New(cfg)
-		flows, err := proto.OnPacket(test.MakePacket(templatePacket), addr)
+		flows, err := proto.OnPacket(test.MakePacket(templatePacket), templateAddr)
 		assert.NoError(t, err)
 		assert.Empty(t, flows)
-
-		v9proto, ok := proto.(*NetflowV9Protocol)
-		assert.True(t, ok)
-		assert.True(t, v9proto.shareTemplates)
-
-		assert.Len(t, v9proto.Session.Sessions, 1)
-		s, found := v9proto.Session.Sessions[key]
-		assert.True(t, found)
-		assert.Len(t, s.Templates, 1)
+		flows, err = proto.OnPacket(test.MakePacket(flowsPacket), flowsAddr)
+		assert.NoError(t, err)
+		assert.Len(t, flows, 1)
 	})
 
 	t.Run("Template sharing disabled", func(t *testing.T) {
-		sharedTemplates := false
-		key := MakeSessionKey(addr, sourceID, sharedTemplates)
 		cfg := config.Defaults()
-		cfg.WithSharedTemplates(sharedTemplates)
+		cfg.WithSharedTemplates(false)
 		proto := New(cfg)
-		flows, err := proto.OnPacket(test.MakePacket(templatePacket), addr)
+		flows, err := proto.OnPacket(test.MakePacket(templatePacket), templateAddr)
+		assert.NoError(t, err)
+		assert.Empty(t, flows)
+		flows, err = proto.OnPacket(test.MakePacket(flowsPacket), flowsAddr)
 		assert.NoError(t, err)
 		assert.Empty(t, flows)
-
-		v9proto, ok := proto.(*NetflowV9Protocol)
-		assert.True(t, ok)
-		assert.False(t, v9proto.shareTemplates)
-
-		assert.Len(t, v9proto.Session.Sessions, 1)
-		s, found := v9proto.Session.Sessions[key]
-		assert.True(t, found)
-		assert.Len(t, s.Templates, 1)
 	})
 }