Add UT for shutdown_config

common/extension/graceful_shutdown.go

@@ -22,8 +22,6 @@ import (
 )
 
 var (
-	// SystemShutdownCallbackNames = []string{"registry"}
-	// systemShutdownCallbacks     = make(map[string]func())
 	customShutdownCallbacks     = list.New()
 )
 

config/graceful_shutdown.go

@@ -116,7 +116,7 @@ func BeforeShutdown() {
 	// If this application is not the consumer, it will do nothing
 	destroyConsumerProtocols()
 
-	logger.Infof("Execute the custom callbacks.")
+	logger.Infof("Graceful shutdown --- Execute the custom callbacks.")
 	customCallbacks := extension.GetAllCustomShutdownCallbacks()
 	for callback := customCallbacks.Front(); callback != nil; callback = callback.Next() {
 		callback.Value.(func())()
@@ -152,12 +152,12 @@ func destroyProviderProtocols() {
 		return
 	}
 
-	consumerProtocol := make(map[string]interface{}, 0)
+	consumerProtocol := make(map[interface{}]interface{}, 0)
 	if consumerConfig != nil && consumerConfig.ProtocolConf != nil {
-		consumerProtocol = consumerConfig.ProtocolConf.(map[string]interface{})
+		consumerProtocol = consumerConfig.ProtocolConf.(map[interface{}]interface{})
 	}
 
-	protocols := providerConfig.ProtocolConf.(map[string]interface{})
+	protocols := providerConfig.ProtocolConf.(map[interface{}]interface{})
 	for name, _ := range protocols {
 		_, found := consumerProtocol[name]
 
@@ -165,7 +165,7 @@ func destroyProviderProtocols() {
 		if found {
 			continue
 		}
-		extension.GetProtocol(name).Destroy()
+		extension.GetProtocol(name.(string)).Destroy()
 	}
 }
 
@@ -202,6 +202,7 @@ func waitForSendingRequests() {
 		// ignore this step
 		return
 	}
+	waitingProcessedTimeout(consumerConfig.ShutdownConfig)
 }
 
 func waitingProcessedTimeout(shutdownConfig *ShutdownConfig) {

config/graceful_shutdown_config_test.go

+/*
+ * Licensed to the Apache Software Foundation (ASF) 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 config
+
+import (
+	"testing"
+	"time"
+)
+import (
+	"github.com/stretchr/testify/assert"
+)
+
+func TestShutdownConfig_GetTimeout(t *testing.T) {
+	config := ShutdownConfig{}
+	assert.False(t, config.RejectRequest)
+	assert.False(t, config.RequestsFinished)
+
+	config = ShutdownConfig{
+		Timeout:     "12x",
+		StepTimeout: "34a",
+	}
+
+	assert.Equal(t, 60*time.Second, config.GetTimeout())
+	assert.Equal(t, 10*time.Second, config.GetStepTimeout())
+
+	config = ShutdownConfig{
+		Timeout:     "34",
+		StepTimeout: "79",
+	}
+
+	assert.Equal(t, 34*time.Millisecond, config.GetTimeout())
+	assert.Equal(t, 79*time.Millisecond, config.GetStepTimeout())
+}

config/graceful_shutdown_test.go

@@ -19,12 +19,16 @@ package config
 
 import (
 	"testing"
-
+)
+import (
 	"github.com/apache/dubbo-go/common/constant"
 	"github.com/apache/dubbo-go/common/extension"
 	"github.com/apache/dubbo-go/protocol"
 )
 
+func TestGracefulShutdownInit(t *testing.T) {
+	GracefulShutdownInit()
+}
 
 func TestBeforeShutdown(t *testing.T) {
 	extension.SetProtocol("registry", func() protocol.Protocol {
@@ -34,23 +38,57 @@ func TestBeforeShutdown(t *testing.T) {
 		return &mockRegistryProtocol{}
 	})
 
+	extension.SetProtocol("mock", func() protocol.Protocol {
+		return &mockRegistryProtocol{}
+	})
+
+
+
 	protocolConfigs := make(map[interface{}]interface{})
 	protocolConfigs[constant.DUBBO] = "aaa"
+
+	// without configuration
+	BeforeShutdown()
+
+	consumerConfig = &ConsumerConfig{
+		ProtocolConf: protocolConfigs,
+		ShutdownConfig: &ShutdownConfig{
+			Timeout:     "1",
+			StepTimeout: "1000",
+		},
+	}
+
+	providerProtocols := make(map[interface{}]interface{})
+	providerProtocols[constant.DUBBO] = "aaa"
+
+	providerProtocols["mock"] = "aaa"
+
 	providerConfig = &ProviderConfig{
 		ShutdownConfig: &ShutdownConfig{
-			Timeout:                       "1",
-			AcceptNewRequestsTimeout:      "1",
-			WaitingProcessRequestsTimeout: "1",
+			Timeout:     "1",
+			StepTimeout: "1000",
+		},
+		ProtocolConf: providerProtocols,
+	}
+	// test destroy protocol
+	BeforeShutdown()
+
+	providerConfig = &ProviderConfig{
+		ShutdownConfig: &ShutdownConfig{
+			Timeout:     "1",
+			StepTimeout: "-1",
 		},
 		ProtocolConf: protocolConfigs,
 	}
+
 	consumerConfig = &ConsumerConfig{
 		ProtocolConf: protocolConfigs,
 		ShutdownConfig: &ShutdownConfig{
-			Timeout:                       "1",
-			AcceptNewRequestsTimeout:      "1",
-			WaitingProcessRequestsTimeout: "1",
+			Timeout:     "1",
+			StepTimeout: "-1",
 		},
 	}
+
+	// test ignore steps
 	BeforeShutdown()
-}
+}
\ No newline at end of file