Fix the test

Signed-off-by: JmPotato <[email protected]>
pingcap · Feb 7, 2023 · 4399c4d · 4399c4d
1 parent 6af6899
commit 4399c4d
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 18 deletions.
diff --git a/domain/infosync/resource_group_manager.go b/domain/infosync/resource_group_manager.go
@@ -23,15 +23,15 @@ import (
 )
 
 type mockResourceGroupManager struct {
-	sync.Mutex
+	sync.RWMutex
 	groups map[string]*rmpb.ResourceGroup
 }
 
 var _ pd.ResourceManagerClient = (*mockResourceGroupManager)(nil)
 
 func (m *mockResourceGroupManager) ListResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, error) {
-	m.Lock()
-	defer m.Unlock()
+	m.RLock()
+	defer m.RUnlock()
 	groups := make([]*rmpb.ResourceGroup, 0, len(m.groups))
 	for _, group := range m.groups {
 		groups = append(groups, group)
@@ -40,8 +40,8 @@ func (m *mockResourceGroupManager) ListResourceGroups(ctx context.Context) ([]*r
 }
 
 func (m *mockResourceGroupManager) GetResourceGroup(ctx context.Context, name string) (*rmpb.ResourceGroup, error) {
-	m.Lock()
-	defer m.Unlock()
+	m.RLock()
+	defer m.RUnlock()
 	group, ok := m.groups[name]
 	if !ok {
 		return nil, nil

diff --git a/store/mockstore/unistore/pd.go b/store/mockstore/unistore/pd.go
@@ -34,17 +34,27 @@ var _ pd.Client = new(pdClient)
 type pdClient struct {
 	*us.MockPD
 
-	serviceSafePoints map[string]uint64
-	gcSafePointMu     sync.Mutex
-	globalConfig      map[string]string
-	externalTimestamp atomic.Uint64
+	serviceSafePoints    map[string]uint64
+	gcSafePointMu        sync.Mutex
+	globalConfig         map[string]string
+	externalTimestamp    atomic.Uint64
+	resourceGroupManager struct {
+		sync.RWMutex
+		groups map[string]*rmpb.ResourceGroup
+	}
 }
 
 func newPDClient(pd *us.MockPD) *pdClient {
 	return &pdClient{
 		MockPD:            pd,
 		serviceSafePoints: make(map[string]uint64),
 		globalConfig:      make(map[string]string),
+		resourceGroupManager: struct {
+			sync.RWMutex
+			groups map[string]*rmpb.ResourceGroup
+		}{
+			groups: make(map[string]*rmpb.ResourceGroup),
+		},
 	}
 }
 
@@ -187,23 +197,44 @@ func (c *pdClient) UpdateKeyspaceState(ctx context.Context, id uint32, state key
 }
 
 func (c *pdClient) ListResourceGroups(ctx context.Context) ([]*rmpb.ResourceGroup, error) {
-	return nil, nil
+	c.resourceGroupManager.RLock()
+	defer c.resourceGroupManager.RUnlock()
+	groups := make([]*rmpb.ResourceGroup, 0, len(c.resourceGroupManager.groups))
+	for _, group := range c.resourceGroupManager.groups {
+		groups = append(groups, group)
+	}
+	return groups, nil
 }
 
-func (c *pdClient) GetResourceGroup(ctx context.Context, resourceGroupName string) (*rmpb.ResourceGroup, error) {
-	return nil, nil
+func (c *pdClient) GetResourceGroup(ctx context.Context, name string) (*rmpb.ResourceGroup, error) {
+	c.resourceGroupManager.RLock()
+	defer c.resourceGroupManager.RUnlock()
+	group, ok := c.resourceGroupManager.groups[name]
+	if !ok {
+		return nil, nil
+	}
+	return group, nil
 }
 
-func (c *pdClient) AddResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
-	return "", nil
+func (c *pdClient) AddResourceGroup(ctx context.Context, group *rmpb.ResourceGroup) (string, error) {
+	c.resourceGroupManager.Lock()
+	defer c.resourceGroupManager.Unlock()
+	c.resourceGroupManager.groups[group.Name] = group
+	return "Success!", nil
 }
 
-func (c *pdClient) ModifyResourceGroup(ctx context.Context, metaGroup *rmpb.ResourceGroup) (string, error) {
-	return "", nil
+func (c *pdClient) ModifyResourceGroup(ctx context.Context, group *rmpb.ResourceGroup) (string, error) {
+	c.resourceGroupManager.Lock()
+	defer c.resourceGroupManager.Unlock()
+	c.resourceGroupManager.groups[group.Name] = group
+	return "Success!", nil
 }
 
-func (c *pdClient) DeleteResourceGroup(ctx context.Context, resourceGroupName string) (string, error) {
-	return "", nil
+func (c *pdClient) DeleteResourceGroup(ctx context.Context, name string) (string, error) {
+	c.resourceGroupManager.Lock()
+	defer c.resourceGroupManager.Unlock()
+	delete(c.resourceGroupManager.groups, name)
+	return "Success!", nil
 }
 
 func (c *pdClient) WatchResourceGroup(ctx context.Context, revision int64) (chan []*rmpb.ResourceGroup, error) {