From 1f0427fe32ec9559a6205e563bc47d44a33ed5e8 Mon Sep 17 00:00:00 2001
From: Michi Mutsuzaki <michi@isovalent.com>
Date: Sun, 27 Aug 2023 00:44:27 +0000
Subject: [PATCH] watcher: Implement fake service functions

Implement FindServiceByIP and add accessor functions to make it easier
to unit test.

Signed-off-by: Michi Mutsuzaki <michi@isovalent.com>
---
 pkg/watcher/fake.go      | 29 ++++++++++++++++--
 pkg/watcher/fake_test.go | 64 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 2 deletions(-)
 create mode 100644 pkg/watcher/fake_test.go

diff --git a/pkg/watcher/fake.go b/pkg/watcher/fake.go
index 73bc565d4f2..cc6cf0e8d58 100644
--- a/pkg/watcher/fake.go
+++ b/pkg/watcher/fake.go
@@ -12,12 +12,18 @@ import (
 // FakeK8sWatcher is used as an "empty" K8sResourceWatcher when --enable-k8s-api flag is not set.
 // It is also used for testing, allowing users to specify a static list of pods.
 type FakeK8sWatcher struct {
-	pods []interface{}
+	pods     []interface{}
+	services []interface{}
 }
 
 // NewK8sWatcher returns a pointer to an initialized FakeK8sWatcher struct.
 func NewFakeK8sWatcher(pods []interface{}) *FakeK8sWatcher {
-	return &FakeK8sWatcher{pods: pods}
+	return NewFakeK8sWatcherWithPodsAndServices(pods, nil)
+}
+
+// NewFakeK8sWatcherWithPodsAndServices returns a pointer to an initialized FakeK8sWatcher struct
+func NewFakeK8sWatcherWithPodsAndServices(pods []interface{}, services []interface{}) *FakeK8sWatcher {
+	return &FakeK8sWatcher{pods, services}
 }
 
 // FindContainer implements K8sResourceWatcher.FindContainer
@@ -34,6 +40,15 @@ func (watcher *FakeK8sWatcher) FindPod(podID string) (*corev1.Pod, error) {
 }
 
 func (watcher *FakeK8sWatcher) FindServiceByIP(ip string) ([]*corev1.Service, error) {
+	for i := range watcher.services {
+		if service, ok := watcher.services[i].(*corev1.Service); ok {
+			for _, serviceIP := range service.Spec.ClusterIPs {
+				if serviceIP == ip {
+					return []*corev1.Service{service}, nil
+				}
+			}
+		}
+	}
 	return nil, fmt.Errorf("service with IP %s not found", ip)
 }
 
@@ -46,3 +61,13 @@ func (watcher *FakeK8sWatcher) AddPod(pod *corev1.Pod) {
 func (watcher *FakeK8sWatcher) ClearAllPods() {
 	watcher.pods = nil
 }
+
+// AddService adds a service to the fake k8s watcher.
+func (watcher *FakeK8sWatcher) AddService(service *corev1.Service) {
+	watcher.services = append(watcher.services, service)
+}
+
+// ClearAllServices removes all services from the fake watcher.
+func (watcher *FakeK8sWatcher) ClearAllServices() {
+	watcher.services = nil
+}
diff --git a/pkg/watcher/fake_test.go b/pkg/watcher/fake_test.go
new file mode 100644
index 00000000000..59457ae3dda
--- /dev/null
+++ b/pkg/watcher/fake_test.go
@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright Authors of Tetragon
+
+package watcher
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func TestFakeK8sWatcher_FindServiceByIP(t *testing.T) {
+	services := []interface{}{
+		&corev1.Service{
+			ObjectMeta: metav1.ObjectMeta{Name: "svc-1"},
+			Spec:       corev1.ServiceSpec{ClusterIPs: []string{"1.1.1.1"}},
+		},
+		&corev1.Service{
+			ObjectMeta: metav1.ObjectMeta{Name: "svc-2"},
+			Spec:       corev1.ServiceSpec{ClusterIPs: []string{"2.2.2.2", "3.3.3.3"}},
+		},
+	}
+	fakeWatcher := NewFakeK8sWatcherWithPodsAndServices(nil, services)
+	res, err := fakeWatcher.FindServiceByIP("1.1.1.1")
+	assert.NoError(t, err)
+	assert.Len(t, res, 1)
+	assert.Equal(t, res[0].Name, "svc-1")
+	res, err = fakeWatcher.FindServiceByIP("2.2.2.2")
+	assert.NoError(t, err)
+	assert.Len(t, res, 1)
+	assert.Equal(t, res[0].Name, "svc-2")
+	res, err = fakeWatcher.FindServiceByIP("3.3.3.3")
+	assert.NoError(t, err)
+	assert.Len(t, res, 1)
+	assert.Equal(t, res[0].Name, "svc-2")
+	_, err = fakeWatcher.FindServiceByIP("4.4.4.4")
+	assert.Error(t, err)
+}
+
+func TestFakeK8sWatcher_AddService(t *testing.T) {
+	services := []interface{}{
+		&corev1.Service{},
+		&corev1.Service{},
+		&corev1.Service{},
+	}
+	fakeWatcher := NewFakeK8sWatcherWithPodsAndServices(nil, services)
+	assert.Len(t, fakeWatcher.services, 3)
+	fakeWatcher.AddService(&corev1.Service{})
+	assert.Len(t, fakeWatcher.services, 4)
+}
+
+func TestFakeK8sWatcher_ClearAllServices(t *testing.T) {
+	services := []interface{}{
+		&corev1.Service{},
+		&corev1.Service{},
+		&corev1.Service{},
+	}
+	fakeWatcher := NewFakeK8sWatcherWithPodsAndServices(nil, services)
+	assert.Len(t, fakeWatcher.services, 3)
+	fakeWatcher.ClearAllServices()
+	assert.Empty(t, fakeWatcher.services)
+}