Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

1064 lines
31 KiB
Go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package telemetry
import (
"context"
"encoding/json"
"fmt"
"sync"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.etcd.io/etcd/api/v3/mvccpb"
clientv3 "go.etcd.io/etcd/client/v3"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
)
// mockKV implements KVInterface for testing
type mockKV struct {
mu sync.Mutex
data map[string]string
shouldFailPut bool
shouldFailGet bool
shouldFailDelete bool
}
func newMockKV() *mockKV {
return &mockKV{
data: make(map[string]string),
}
}
func (m *mockKV) Put(ctx context.Context, key, val string) error {
if m.shouldFailPut {
return errors.New("mock put error")
}
m.mu.Lock()
defer m.mu.Unlock()
m.data[key] = val
return nil
}
func (m *mockKV) Get(ctx context.Context, key string, opts ...clientv3.OpOption) (*clientv3.GetResponse, error) {
if m.shouldFailGet {
return nil, errors.New("mock get error")
}
m.mu.Lock()
defer m.mu.Unlock()
var kvs []*mvccpb.KeyValue
// Check if it's a prefix query by looking at the key
for k, v := range m.data {
if len(k) >= len(key) && k[:len(key)] == key {
kvs = append(kvs, &mvccpb.KeyValue{
Key: []byte(k),
Value: []byte(v),
})
}
}
return &clientv3.GetResponse{
Kvs: kvs,
}, nil
}
func (m *mockKV) Delete(ctx context.Context, key string, opts ...clientv3.OpOption) error {
if m.shouldFailDelete {
return errors.New("mock delete error")
}
m.mu.Lock()
defer m.mu.Unlock()
delete(m.data, key)
return nil
}
// =============================================================================
// Core Functionality Tests
// =============================================================================
func TestNewCommandStore(t *testing.T) {
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/prefix/")
assert.NotNil(t, store)
assert.Equal(t, "/test/prefix/configs/", store.configPath)
assert.NotNil(t, store.cache)
assert.Empty(t, store.cache.commands)
assert.Empty(t, store.cache.configs)
}
// =============================================================================
// Command CRUD Tests (In-Memory Only)
// =============================================================================
func TestCommandPushAndRetrieve(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push a command
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "show_errors",
Payload: []byte(`{"max_count": 50}`),
})
require.NoError(t, err)
assert.NotEmpty(t, cmdID)
// Commands are NOT stored in etcd (in-memory only)
assert.Empty(t, mockKv.data, "commands should not be persisted to etcd")
// Retrieve commands
commands, err := store.ListCommands(ctx)
require.NoError(t, err)
assert.Len(t, commands, 1)
assert.Equal(t, cmdID, commands[0].CommandId)
assert.Equal(t, "show_errors", commands[0].CommandType)
assert.Equal(t, []byte(`{"max_count": 50}`), commands[0].Payload)
assert.Equal(t, "global", commands[0].TargetScope)
}
func TestCommandWithTargetClient(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "collection_metrics",
TargetClientId: "client-abc",
Payload: []byte(`{"collections": ["col1"]}`),
})
require.NoError(t, err)
commands, err := store.ListCommands(ctx)
require.NoError(t, err)
assert.Len(t, commands, 1)
assert.Equal(t, cmdID, commands[0].CommandId)
assert.Equal(t, "client:client-abc", commands[0].TargetScope)
}
func TestCommandWithTargetDatabase(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "update_setting",
TargetDatabase: "test_db",
Payload: []byte(`{"setting": "value"}`),
})
require.NoError(t, err)
commands, err := store.ListCommands(ctx)
require.NoError(t, err)
assert.Len(t, commands, 1)
assert.Equal(t, cmdID, commands[0].CommandId)
assert.Equal(t, "database:test_db", commands[0].TargetScope)
}
func TestCommandWithTargetClientTakesPrecedenceOverDatabase(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// When both TargetClientId and TargetDatabase are set, TargetClientId takes precedence
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "update_setting",
TargetClientId: "client-xyz",
TargetDatabase: "test_db",
Payload: []byte(`{"setting": "value"}`),
})
require.NoError(t, err)
commands, err := store.ListCommands(ctx)
require.NoError(t, err)
assert.Len(t, commands, 1)
assert.Equal(t, cmdID, commands[0].CommandId)
assert.Equal(t, "client:client-xyz", commands[0].TargetScope)
}
func TestCommandDelete(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push a command
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "show_errors",
})
require.NoError(t, err)
// Verify it exists
commands, err := store.ListCommands(ctx)
require.NoError(t, err)
assert.Len(t, commands, 1)
// Delete it
err = store.DeleteCommand(ctx, cmdID)
require.NoError(t, err)
// Verify it's gone
commands, err = store.ListCommands(ctx)
require.NoError(t, err)
assert.Empty(t, commands)
}
func TestCommandTTLExpiration(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push a command with 1 second TTL
expiredID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "expired_cmd",
TtlSeconds: 1,
})
require.NoError(t, err)
// Manually set create time to 2 seconds ago (expired)
store.cacheMu.Lock()
store.cache.commands[expiredID].CreateTime = time.Now().Add(-2 * time.Second).UnixMilli()
store.cacheMu.Unlock()
// Push a valid command (no TTL)
validID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "valid_cmd",
TtlSeconds: 0,
})
require.NoError(t, err)
// ListCommands should filter out expired command
commands, err := store.ListCommands(ctx)
require.NoError(t, err)
assert.Len(t, commands, 1)
assert.Equal(t, validID, commands[0].CommandId)
}
func TestCleanupExpiredCommands(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push an expired command
expiredID, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "expired",
TtlSeconds: 1,
})
store.cacheMu.Lock()
store.cache.commands[expiredID].CreateTime = time.Now().Add(-2 * time.Second).UnixMilli()
store.cacheMu.Unlock()
// Push a valid command
validID, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "valid",
TtlSeconds: 3600,
})
// Run cleanup
store.CleanupExpiredCommands(ctx)
// Verify only valid command remains
store.cacheMu.RLock()
_, expiredExists := store.cache.commands[expiredID]
_, validExists := store.cache.commands[validID]
store.cacheMu.RUnlock()
assert.False(t, expiredExists, "expired command should be removed")
assert.True(t, validExists, "valid command should remain")
}
// =============================================================================
// Config CRUD Tests (Persisted to etcd)
// =============================================================================
func TestConfigPushAndRetrieve(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push a persistent config
cfgID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"enabled": true}`),
Persistent: true,
})
require.NoError(t, err)
assert.NotEmpty(t, cfgID)
// Configs ARE stored in etcd
assert.Len(t, mockKv.data, 1, "configs should be persisted to etcd")
// Retrieve configs
configs, hash, err := store.ListConfigs(ctx)
require.NoError(t, err)
assert.Len(t, configs, 1)
assert.Equal(t, cfgID, configs[0].ConfigId)
assert.Equal(t, "push_config", configs[0].ConfigType)
assert.NotEmpty(t, hash)
}
func TestConfigDelete(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push a config
cfgID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Persistent: true,
})
require.NoError(t, err)
assert.Len(t, mockKv.data, 1)
// Delete it
err = store.DeleteCommand(ctx, cfgID)
require.NoError(t, err)
// Verify removed from both cache and etcd
configs, _, err := store.ListConfigs(ctx)
require.NoError(t, err)
assert.Empty(t, configs)
assert.Empty(t, mockKv.data)
}
func TestOnlyPushConfigCanBePersistent(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// push_config can be persistent
_, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Persistent: true,
})
assert.NoError(t, err)
// Other command types cannot be persistent
invalidTypes := []string{"show_errors", "collection_metrics", "custom_command"}
for _, cmdType := range invalidTypes {
_, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: cmdType,
Persistent: true,
})
assert.Error(t, err, "command type %s should not be allowed to be persistent", cmdType)
assert.Contains(t, err.Error(), "only push_config can be persistent")
}
}
// =============================================================================
// Config Hash Tests
// =============================================================================
func TestConfigHashConsistency(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Empty configs should have empty hash
_, hash1, err := store.ListConfigs(ctx)
require.NoError(t, err)
assert.Empty(t, hash1)
// Add a config for client1 (scope: client:client1)
cfgID1, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"a": 1}`),
Persistent: true,
TargetClientId: "client1",
})
_, hash2, err := store.ListConfigs(ctx)
require.NoError(t, err)
assert.NotEmpty(t, hash2)
assert.Len(t, hash2, 16)
// Same configs should produce same hash
_, hash3, _ := store.ListConfigs(ctx)
assert.Equal(t, hash2, hash3)
// Add another config for client2 (different scope) - hash should change
cfgID2, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"b": 2}`),
Persistent: true,
TargetClientId: "client2",
})
_, hash4, _ := store.ListConfigs(ctx)
assert.NotEqual(t, hash2, hash4, "hash should change when config added")
// Delete a config - hash should change again
store.DeleteCommand(ctx, cfgID1)
_, hash5, _ := store.ListConfigs(ctx)
assert.NotEqual(t, hash4, hash5, "hash should change when config deleted")
// Cleanup
store.DeleteCommand(ctx, cfgID2)
}
// =============================================================================
// Restart Behavior Test (Critical!)
// =============================================================================
func TestRestartBehavior_CommandsLostConfigsPersist(t *testing.T) {
ctx := context.Background()
// Use shared mockKV to simulate etcd persistence across restarts
sharedKV := newMockKV()
// === First "session" before restart ===
store1 := NewCommandStoreWithKV(sharedKV, "/test/")
// Push a non-persistent command (in-memory only)
cmdID, err := store1.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "show_errors",
Payload: []byte(`{"max_count": 100}`),
})
require.NoError(t, err)
// Push a persistent config (stored in etcd)
cfgID, err := store1.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"enabled": true, "sampling_rate": 0.5}`),
Persistent: true,
})
require.NoError(t, err)
// Verify both exist before restart
commands1, _ := store1.ListCommands(ctx)
assert.Len(t, commands1, 1, "command should exist before restart")
assert.Equal(t, cmdID, commands1[0].CommandId)
configs1, hash1, _ := store1.ListConfigs(ctx)
assert.Len(t, configs1, 1, "config should exist before restart")
assert.Equal(t, cfgID, configs1[0].ConfigId)
assert.NotEmpty(t, hash1)
// Verify etcd state: only config is persisted
assert.Len(t, sharedKV.data, 1, "only config should be in etcd")
// === Simulate restart by creating new store with same etcd ===
store2 := NewCommandStoreWithKV(sharedKV, "/test/")
// Commands should be GONE (they were in-memory only)
commands2, err := store2.ListCommands(ctx)
require.NoError(t, err)
assert.Empty(t, commands2, "commands should be lost after restart")
// Configs should PERSIST (loaded from etcd)
configs2, hash2, err := store2.ListConfigs(ctx)
require.NoError(t, err)
assert.Len(t, configs2, 1, "configs should persist after restart")
assert.Equal(t, cfgID, configs2[0].ConfigId)
assert.Equal(t, "push_config", configs2[0].ConfigType)
assert.Equal(t, []byte(`{"enabled": true, "sampling_rate": 0.5}`), configs2[0].Payload)
// Hash should be the same after restart
assert.Equal(t, hash1, hash2, "config hash should be consistent after restart")
}
func TestRestartBehavior_MultipleConfigs(t *testing.T) {
ctx := context.Background()
sharedKV := newMockKV()
// First session: create multiple configs with different scopes
store1 := NewCommandStoreWithKV(sharedKV, "/test/")
cfgID1, _ := store1.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"setting": "value1"}`),
Persistent: true,
TargetClientId: "client1",
})
cfgID2, _ := store1.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"setting": "value2"}`),
Persistent: true,
TargetClientId: "client2",
})
_, hash1, _ := store1.ListConfigs(ctx)
// Simulate restart
store2 := NewCommandStoreWithKV(sharedKV, "/test/")
configs2, hash2, err := store2.ListConfigs(ctx)
require.NoError(t, err)
assert.Len(t, configs2, 2, "all configs should persist")
// Verify both configs are present
configIDs := make(map[string]bool)
for _, cfg := range configs2 {
configIDs[cfg.ConfigId] = true
}
assert.True(t, configIDs[cfgID1])
assert.True(t, configIDs[cfgID2])
// Hash should be consistent
assert.Equal(t, hash1, hash2)
}
// =============================================================================
// ListCommandsWithInfo Tests
// =============================================================================
func TestListCommandsWithInfo(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push a command
cmdID, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "show_errors",
TtlSeconds: 3600,
})
// Push a config
cfgID, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Persistent: true,
})
// List all with info
infos, err := store.ListCommandsWithInfo(ctx)
require.NoError(t, err)
assert.Len(t, infos, 2)
// Verify command info
var cmdInfo, cfgInfo *CommandInfoData
for _, info := range infos {
switch info.CommandID {
case cmdID:
cmdInfo = info
case cfgID:
cfgInfo = info
}
}
require.NotNil(t, cmdInfo)
assert.Equal(t, "show_errors", cmdInfo.CommandType)
assert.False(t, cmdInfo.Persistent)
assert.Equal(t, int64(3600), cmdInfo.TTLSeconds)
require.NotNil(t, cfgInfo)
assert.Equal(t, "push_config", cfgInfo.CommandType)
assert.True(t, cfgInfo.Persistent)
assert.Equal(t, int64(0), cfgInfo.TTLSeconds)
}
// =============================================================================
// Error Handling Tests
// =============================================================================
func TestPushConfigError(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
mockKv.shouldFailPut = true
store := NewCommandStoreWithKV(mockKv, "/test/")
_, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Persistent: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "IO failed")
}
func TestDeleteConfigError(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
cfgID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Persistent: true,
})
require.NoError(t, err)
mockKv.shouldFailDelete = true
err = store.DeleteCommand(ctx, cfgID)
assert.Error(t, err)
assert.Contains(t, err.Error(), "IO failed")
}
func TestLoadCacheWithFailedGet(t *testing.T) {
mockKv := newMockKV()
mockKv.shouldFailGet = true
// Should not panic, just log warning and return empty cache
store := NewCommandStoreWithKV(mockKv, "/test/")
assert.NotNil(t, store)
assert.Empty(t, store.cache.configs)
}
// =============================================================================
// Concurrent Access Tests
// =============================================================================
func TestConcurrentAccess(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push some initial data
for i := 0; i < 5; i++ {
store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "show_errors",
})
}
// Concurrent reads
const numGoroutines = 20
var wg sync.WaitGroup
wg.Add(numGoroutines)
for i := 0; i < numGoroutines; i++ {
go func() {
defer wg.Done()
commands, err := store.ListCommands(ctx)
assert.NoError(t, err)
assert.Len(t, commands, 5)
}()
}
wg.Wait()
}
// =============================================================================
// Payload Preservation Tests
// =============================================================================
func TestPayloadPreservation(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
testPayloads := [][]byte{
[]byte(`{"key": "value"}`),
[]byte(`{"unicode": "✓✗✔"}`),
[]byte(`{"quotes": "\"quoted\""}`),
{0, 1, 2, 3, 255}, // binary data
nil, // nil payload
[]byte(""), // empty payload
}
// Use different TargetClientId for each to create separate configs
for i, payload := range testPayloads {
_, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: payload,
Persistent: true,
TargetClientId: fmt.Sprintf("client%d", i),
})
require.NoError(t, err, "failed for payload %d", i)
}
configs, _, err := store.ListConfigs(ctx)
require.NoError(t, err)
assert.Len(t, configs, len(testPayloads))
}
// =============================================================================
// Cache Behavior Tests
// =============================================================================
func TestCacheNotAffectedByDirectEtcdChanges(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push a config through store
cfgID, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"a": 1}`),
Persistent: true,
})
configs1, hash1, _ := store.ListConfigs(ctx)
assert.Len(t, configs1, 1)
// Directly modify etcd (bypassing store/cache)
newCfg := storedConfig{
ConfigID: "cfg-fake",
ConfigType: "push_config",
Payload: []byte(`{"fake": 1}`),
CreateTime: time.Now().UnixMilli(),
TargetScope: "global",
}
data, _ := json.Marshal(newCfg)
mockKv.data["/test/configs/cfg-fake"] = string(data)
// Store should NOT see direct etcd changes (cache is authoritative)
configs2, hash2, _ := store.ListConfigs(ctx)
assert.Len(t, configs2, 1, "direct etcd changes should not be visible")
assert.Equal(t, hash1, hash2)
// Cleanup
store.DeleteCommand(ctx, cfgID)
}
// =============================================================================
// DeleteNonPersistentCommand Tests
// =============================================================================
func TestDeleteNonPersistentCommand(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
t.Run("deletes existing non-persistent command", func(t *testing.T) {
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "test_cmd",
Payload: []byte(`{"test": true}`),
Persistent: false,
})
require.NoError(t, err)
// Verify command exists
commands, _ := store.ListCommands(ctx)
assert.Len(t, commands, 1)
// Delete the command
deleted := store.DeleteNonPersistentCommand(cmdID)
assert.True(t, deleted)
// Verify command is gone
commands, _ = store.ListCommands(ctx)
assert.Empty(t, commands)
})
t.Run("returns false for non-existent command", func(t *testing.T) {
deleted := store.DeleteNonPersistentCommand("non-existent-id")
assert.False(t, deleted)
})
t.Run("does not delete persistent configs", func(t *testing.T) {
cfgID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"config": true}`),
Persistent: true,
})
require.NoError(t, err)
// Try to delete - should return false (config not in commands map)
deleted := store.DeleteNonPersistentCommand(cfgID)
assert.False(t, deleted)
// Config should still exist
configs, _, _ := store.ListConfigs(ctx)
assert.Len(t, configs, 1)
// Cleanup
store.DeleteCommand(ctx, cfgID)
})
}
// =============================================================================
// GetCommandInfo Tests
// =============================================================================
func TestGetCommandInfo(t *testing.T) {
ctx := context.Background()
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
t.Run("returns info for non-persistent command", func(t *testing.T) {
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "show_errors",
Payload: []byte(`{"max_count": 10}`),
Persistent: false,
})
require.NoError(t, err)
cmdType, payload, persistent, ok := store.GetCommandInfo(cmdID)
assert.True(t, ok)
assert.Equal(t, "show_errors", cmdType)
assert.Equal(t, []byte(`{"max_count": 10}`), payload)
assert.False(t, persistent)
// Cleanup
store.DeleteNonPersistentCommand(cmdID)
})
t.Run("returns info for persistent config", func(t *testing.T) {
cfgID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"setting": "value"}`),
Persistent: true,
})
require.NoError(t, err)
cmdType, payload, persistent, ok := store.GetCommandInfo(cfgID)
assert.True(t, ok)
assert.Equal(t, "push_config", cmdType)
assert.Equal(t, []byte(`{"setting": "value"}`), payload)
assert.True(t, persistent)
// Cleanup
store.DeleteCommand(ctx, cfgID)
})
t.Run("returns false for non-existent ID", func(t *testing.T) {
_, _, _, ok := store.GetCommandInfo("non-existent")
assert.False(t, ok)
})
}
// =============================================================================
// mergeJSONPayloads Tests
// =============================================================================
func TestMergeJSONPayloads(t *testing.T) {
t.Run("returns false for empty new payload", func(t *testing.T) {
_, ok := mergeJSONPayloads([][]byte{[]byte(`{"a": 1}`)}, []byte{})
assert.False(t, ok)
_, ok = mergeJSONPayloads([][]byte{[]byte(`{"a": 1}`)}, nil)
assert.False(t, ok)
})
t.Run("returns false for invalid new payload JSON", func(t *testing.T) {
_, ok := mergeJSONPayloads([][]byte{[]byte(`{"a": 1}`)}, []byte(`invalid json`))
assert.False(t, ok)
})
t.Run("merges valid payloads", func(t *testing.T) {
existing := [][]byte{
[]byte(`{"a": 1, "b": 2}`),
[]byte(`{"c": 3}`),
}
newPayload := []byte(`{"d": 4, "a": 10}`) // a is overwritten
merged, ok := mergeJSONPayloads(existing, newPayload)
assert.True(t, ok)
var result map[string]interface{}
err := json.Unmarshal(merged, &result)
require.NoError(t, err)
assert.Equal(t, float64(10), result["a"]) // overwritten
assert.Equal(t, float64(2), result["b"])
assert.Equal(t, float64(3), result["c"])
assert.Equal(t, float64(4), result["d"])
})
t.Run("skips invalid existing payloads", func(t *testing.T) {
existing := [][]byte{
[]byte(`{"a": 1}`),
[]byte(`invalid json`), // should be skipped
[]byte(`{"b": 2}`),
}
newPayload := []byte(`{"c": 3}`)
merged, ok := mergeJSONPayloads(existing, newPayload)
assert.True(t, ok)
var result map[string]interface{}
err := json.Unmarshal(merged, &result)
require.NoError(t, err)
assert.Equal(t, float64(1), result["a"])
assert.Equal(t, float64(2), result["b"])
assert.Equal(t, float64(3), result["c"])
})
t.Run("merges with empty existing", func(t *testing.T) {
merged, ok := mergeJSONPayloads([][]byte{}, []byte(`{"x": 100}`))
assert.True(t, ok)
var result map[string]interface{}
err := json.Unmarshal(merged, &result)
require.NoError(t, err)
assert.Equal(t, float64(100), result["x"])
})
}
// =============================================================================
// PushCommand Edge Cases
// =============================================================================
func TestPushCommandEdgeCases(t *testing.T) {
ctx := context.Background()
t.Run("merges existing configs with same scope", func(t *testing.T) {
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Push first config
cfg1ID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"setting1": "value1"}`),
Persistent: true,
})
require.NoError(t, err)
// Push second config with same scope - should replace first
cfg2ID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"setting2": "value2"}`),
Persistent: true,
})
require.NoError(t, err)
// First config should be deleted, only second should exist
configs, _, _ := store.ListConfigs(ctx)
assert.Len(t, configs, 1)
assert.Equal(t, cfg2ID, configs[0].ConfigId)
assert.NotEqual(t, cfg1ID, cfg2ID)
})
t.Run("handles put error for config", func(t *testing.T) {
mockKv := newMockKV()
mockKv.shouldFailPut = true
store := NewCommandStoreWithKV(mockKv, "/test/")
_, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Payload: []byte(`{"test": true}`),
Persistent: true,
})
assert.Error(t, err)
})
t.Run("handles nil payload", func(t *testing.T) {
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
cmdID, err := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "test_cmd",
Payload: nil,
Persistent: false,
})
require.NoError(t, err)
assert.NotEmpty(t, cmdID)
// Cleanup
store.DeleteNonPersistentCommand(cmdID)
})
}
// =============================================================================
// LoadCache Edge Cases
// =============================================================================
func TestLoadCacheEdgeCases(t *testing.T) {
t.Run("handles get error during load", func(t *testing.T) {
mockKv := newMockKV()
mockKv.shouldFailGet = true
// Should not panic, just skip loading
store := NewCommandStoreWithKV(mockKv, "/test/")
assert.NotNil(t, store)
// Cache should be empty
configs, _, _ := store.ListConfigs(context.Background())
assert.Empty(t, configs)
})
t.Run("handles invalid JSON in etcd", func(t *testing.T) {
mockKv := newMockKV()
// Pre-populate with invalid JSON
mockKv.data["/test/configs/bad-config"] = "not valid json"
// Should not panic, just skip invalid entries
store := NewCommandStoreWithKV(mockKv, "/test/")
assert.NotNil(t, store)
// Cache should be empty (invalid entry skipped)
configs, _, _ := store.ListConfigs(context.Background())
assert.Empty(t, configs)
})
}
// =============================================================================
// ListCommandsWithInfo Edge Cases
// =============================================================================
func TestListCommandsWithInfoEdgeCases(t *testing.T) {
ctx := context.Background()
t.Run("includes both commands and configs", func(t *testing.T) {
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Add a non-persistent command
cmdID, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "one_time_cmd",
Persistent: false,
TtlSeconds: 3600,
})
// Add a persistent config
cfgID, _ := store.PushCommand(ctx, &milvuspb.PushClientCommandRequest{
CommandType: "push_config",
Persistent: true,
})
infos, err := store.ListCommandsWithInfo(ctx)
require.NoError(t, err)
assert.Len(t, infos, 2)
// Find command and config in results
var foundCmd, foundCfg bool
for _, info := range infos {
if info.CommandID == cmdID {
foundCmd = true
assert.False(t, info.Persistent)
assert.Equal(t, int64(3600), info.TTLSeconds)
}
if info.CommandID == cfgID {
foundCfg = true
assert.True(t, info.Persistent)
assert.Equal(t, int64(0), info.TTLSeconds)
}
}
assert.True(t, foundCmd, "command not found")
assert.True(t, foundCfg, "config not found")
// Cleanup
store.DeleteNonPersistentCommand(cmdID)
store.DeleteCommand(ctx, cfgID)
})
t.Run("skips expired commands", func(t *testing.T) {
mockKv := newMockKV()
store := NewCommandStoreWithKV(mockKv, "/test/")
// Manually add an expired command to the cache
store.cacheMu.Lock()
store.cache.commands["expired-cmd"] = &storedCommand{
CommandID: "expired-cmd",
CommandType: "old_cmd",
CreateTime: time.Now().Add(-2 * time.Hour).UnixMilli(),
TTLSeconds: 1, // 1 second TTL, already expired
}
store.cacheMu.Unlock()
infos, err := store.ListCommandsWithInfo(ctx)
require.NoError(t, err)
// Expired command should be skipped
for _, info := range infos {
assert.NotEqual(t, "expired-cmd", info.CommandID)
}
})
}