Repository: DiceDB/dice
File: integration_tests/commands/resp/pfcountwatch_test.go (Go)

Code before:

	channels := make([]<-chan *dicedb.WatchResult, len(subscribers))
	for i, subscriber := range subscribers {
		watch := subscriber.client.WatchConn(context.Background())
		subscribers[i].watch = watch
		assert.Assert(t, watch != nil)
		firstMsg, err := watch.Watch(context.Background(), "PFCOUNT", pfcountWatchKey)
		assert.NilError(t, err)
		assert.Equal(t, firstMsg.Command, "PFCOUNT")
		channels[i] = watch.Channel()
	}

	for _, tc := range PFCountWatchSDKTestCases {
		err := publisher.PFAdd(context.Background(), tc.key, tc.val).Err()
		assert.NilError(t, err)

		for _, channel := range channels {
			v := <-channel

			assert.Equal(t, "PFCOUNT", v.Command)         // command
			assert.Equal(t, "1580567186", v.Fingerprint) // Fingerprint
			assert.DeepEqual(t, tc.result, v.Data)       // data
		}
	}

	for _, tc := range pfcountWatchSDKhWithPFMergeTestCases {
		publisher.Del(context.Background(), tc.dest_key_1, tc.dest_key_2)
		publisher.PFAdd(context.Background(), tc.dest_key_1, tc.dest_value_1).Err()
		publisher.PFAdd(context.Background(), tc.dest_key_2, tc.dest_value_2).Err()
		publisher.PFMerge(context.Background(), pfcountWatchKey, tc.dest_key_1, tc.dest_key_2).Err()

		for _, channel := range channels {
			v := <-channel

			assert.Equal(t, "PFCOUNT", v.Command)         // command
			assert.Equal(t, "1580567186", v.Fingerprint) // Fingerprint
			assert.DeepEqual(t, tc.result, v.Data)       // data
		}
	}
}

Review comment (refactor):
Maybe we can rewrite this better as follows:
```go
func TestPFCountWatchWithSDK(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
	defer cancel()

	publisher := getLocalSDK(t)
	subscribers := setupSubscribers(t, 3)
	defer cleanupSubscribers(subscribers)

	// Clean up test key before starting
	err := publisher.Del(ctx, pfCountWatchKey)
	assert.NilError(t, err)

	// Setup watch channels
	channels := setupWatchChannels(t, ctx, subscribers)

	t.Run("Basic PFCount Operations", func(t *testing.T) {
		testPFCountOperations(t, ctx, publisher, channels)
	})

	t.Run("PFCount with Merge Operations", func(t *testing.T) {
		testPFCountMergeOperations(t, ctx, publisher, channels)
	})
}

func setupSubscribers(t *testing.T, count int) []watchSubscriber {
	subscribers := make([]watchSubscriber, count)
	for i := range subscribers {
		subscribers[i].client = getLocalSDK(t)
	}
	return subscribers
}

func cleanupSubscribers(subscribers []watchSubscriber) {
	for _, sub := range subscribers {
		if sub.watch != nil {
			sub.watch.Close()
		}
	}
}

func setupWatchChannels(t *testing.T, ctx context.Context, subscribers []watchSubscriber) []<-chan *dicedb.WatchResult {
	channels := make([]<-chan *dicedb.WatchResult, len(subscribers))
	
	for i, subscriber := range subscribers {
		watch := subscriber.client.WatchConn(ctx)
		assert.Assert(t, watch != nil)
		
		firstMsg, err := watch.Watch(ctx, "PFCOUNT", pfCountWatchKey)
		assert.NilError(t, err)
		assert.Equal(t, firstMsg.Command, "PFCOUNT")
		
		subscribers[i].watch = watch
		channels[i] = watch.Channel()
	}
	
	return channels
}

func testPFCountOperations(t *testing.T, ctx context.Context, publisher *SDK, channels []<-chan *dicedb.WatchResult) {
	for _, tc := range pfCountTestCases {
		err := publisher.PFAdd(ctx, tc.Key, tc.Value)
		assert.NilError(t, err)

		verifyWatchResults(t, channels, tc.Result)
	}
}

func testPFCountMergeOperations(t *testing.T, ctx context.Context, publisher *SDK, channels []<-chan *dicedb.WatchResult) {
	for _, tc := range pfCountMergeTestCases {
		// Clean up previous keys
		err := publisher.Del(ctx, tc.DestKey1, tc.DestKey2)
		assert.NilError(t, err)

		// Add values to both keys
		err = publisher.PFAdd(ctx, tc.DestKey1, tc.DestValue1...)
		assert.NilError(t, err)
		err = publisher.PFAdd(ctx, tc.DestKey2, tc.DestValue2...)
		assert.NilError(t, err)

		// Merge keys
		err = publisher.PFMerge(ctx, pfCountWatchKey, tc.DestKey1, tc.DestKey2)
		assert.NilError(t, err)

		verifyWatchResults(t, channels, tc.Expected)
	}
}

func verifyWatchResults(t *testing.T, channels []<-chan *dicedb.WatchResult, expected int64) {
	for _, channel := range channels {
		select {
		case v := <-channel:
			assert.Equal(t, "PFCOUNT", v.Command)
			assert.Equal(t, watchFingerprint, v.Fingerprint)
			assert.DeepEqual(t, expected, v.Data)
		case <-time.After(defaultTimeout):
			t.Fatal("timeout waiting for watch result")
		}
	}
}
```

Code after:
	val    any
	result int64
}

type pfcountWatchSDKWithPFMergeTestCase struct {
	destKey1    string
	destValue1  []string
	destKey2 	string
	destValue2  []string
	result 		int64
}

var PFCountWatchSDKTestCases = []pfcountWatchSDKTestCase{
	{pfcountWatchKeySDK, "value1", int64(1)},
	{pfcountWatchKeySDK, "value2", int64(2)},
	{pfcountWatchKeySDK, "value3", int64(3)},
	{pfcountWatchKeySDK, "value4", int64(4)},
}

var pfcountWatchSDKhWithPFMergeTestCases = []pfcountWatchSDKWithPFMergeTestCase{
	{"DEST_KEY_1", []string{"a", "b", "c", "d", "e", "f"}, "DEST_KEY_2", []string{"f", "g"}, int64(11)},
	{"DEST_KEY_3", []string{"a", "b", "c", "d", "e", "f"}, "DEST_KEY_4", []string{"a1", "b1", "c1", "d1", "e1", "f1"}, int64(17)},
	{"DEST_KEY_5", []string{"f", "g"}, "DEST_KEY_6", []string{"f", "g"}, int64(17)},
}

func TestPFCountWATCHWithSDK(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout) 
	defer cancel()

	publisher := getLocalSdk()
	subscribers := setupSubscribersSDK(3)
	defer cleanupSubscribersSDK(subscribers)

	publisher.Del(ctx, pfcountWatchKey)

	channels := setUpWatchChannelsSDK(t, ctx, subscribers)

	t.Run("Basic PFCount Operations", func(t *testing.T) { 
			testPFCountAddSDK(t, ctx, channels, publisher) 
		},
	)

	t.Run("PFCount Operations including PFMerge", func(t *testing.T) { 
			testPFCountMergeSDK(t, ctx, channels, publisher) 
		},
	)
}

func setupSubscribersSDK(count int) []WatchSubscriber { 
	subscribers := make([]WatchSubscriber, count) 
	for i := range subscribers {