making a pre-test service

Author: pjohnst5

cns/restserver/api_test.go

@@ -176,9 +176,10 @@ func TestMain(m *testing.M) {
 	logger.InitLogger("testlogs", 0, 0, "./")
 
 	// Create the service. If CRD channel mode is needed, then at the start of the test,
-	// it can stop the service (service.Stop), invoke startService again with new ServiceConfig (with CRD mode)
-	// perform the test and then restore the service again.
-	if err = startService(common.ServiceConfig{ChannelMode: cns.Direct}, configuration.CNSConfig{}); err != nil {
+	// it can stop the service (via the returned stop func), invoke startService again with
+	// new ServiceConfig (with CRD mode), perform the test and then restore the service again.
+	stop, _, err := startService(common.ServiceConfig{ChannelMode: cns.Direct}, configuration.CNSConfig{})
+	if err != nil {
 		fmt.Printf("Failed to start CNS Service. Error: %v", err)
 		os.Exit(1)
 	}
@@ -207,7 +208,8 @@ func TestMain(m *testing.M) {
 	exitCode := m.Run()
 
 	// Cleanup.
-	service.Stop()
+	// Use the stop func returned by startService to ensure proper teardown.
+	stop()
 	nmAgentServer.Stop()
 
 	os.Exit(exitCode)
@@ -1031,7 +1033,7 @@ func publishNCViaCNS(
 
 		if bodyStatus != expStatus {
 			// nolint:goerr113 // this is okay in a test:
-			return fmt.Errorf("unexpected status in body. exp: %d, got %d", expStatus, bodyStatus)
+			return fmt.Errorf("unexpected status in body. exp: %d, got: %d", expStatus, bodyStatus)
 		}
 	}
 
@@ -1787,23 +1789,23 @@ func setEnv(t *testing.T) *httptest.ResponseRecorder {
 	return w
 }
 
-func startService(serviceConfig common.ServiceConfig, _ configuration.CNSConfig) error {
+func startService(serviceConfig common.ServiceConfig, _ configuration.CNSConfig) (func(), *HTTPRestService, error) {
 	// Create the service.
 	config := serviceConfig
 
 	// Create the key value fileStore.
 	fileStore, err := store.NewJsonFileStore(cnsJsonFileName, processlock.NewMockFileLock(false), nil)
 	if err != nil {
 		logger.Errorf("Failed to create store file: %s, due to error %v\n", cnsJsonFileName, err)
-		return err
+		return nil, nil, err
 	}
 	config.Store = fileStore
 
 	nmagentClient := &fakes.NMAgentClientFake{}
 	service, err = NewHTTPRestService(&config, &fakes.WireserverClientFake{}, &fakes.WireserverProxyFake{}, &IPtablesProvider{},
 		nmagentClient, nil, nil, nil, fakes.NewMockIMDSClient())
 	if err != nil {
-		return err
+		return nil, nil, err
 	}
 	svc = service.(*HTTPRestService)
 	svc.Service.Options[acncommon.OptCnsURL] = ""
@@ -1839,6 +1841,7 @@ func startService(serviceConfig common.ServiceConfig, _ configuration.CNSConfig)
 		return resp, nil
 	}
 
+	var perTestSvc *HTTPRestService
 	if service != nil {
 		// Create empty azure-cns.json. CNS should start successfully by deleting this file
 		file, _ := os.Create(cnsJsonFileName)
@@ -1850,25 +1853,43 @@ func startService(serviceConfig common.ServiceConfig, _ configuration.CNSConfig)
 		err = service.Init(&config)
 		if err != nil {
 			logger.Errorf("Failed to Init CNS, err:%v.\n", err)
-			return err
+			return nil, nil, err
 		}
 
 		err = service.Start(&config)
 		if err != nil {
 			logger.Errorf("Failed to start CNS, err:%v.\n", err)
-			return err
+			return nil, nil, err
 		}
 
 		if _, err := os.Stat(cnsJsonFileName); err == nil || !os.IsNotExist(err) {
 			logger.Errorf("Failed to remove empty CNS state file: %s, err:%v", cnsJsonFileName, err)
-			return err
+			return nil, nil, err
 		}
+		perTestSvc = service.(*HTTPRestService)
 	}
 
 	// Get the internal http mux as test hook.
-	mux = service.(*HTTPRestService).Listener.GetMux()
+	if perTestSvc == nil {
+		return nil, nil, fmt.Errorf("service was not created")
+	}
+	mux = perTestSvc.Listener.GetMux()
 
-	return nil
+	// stop func cleans up the running service and any test artifacts.
+	stop := func() {
+		if service != nil {
+			// ignore stop errors
+			service.Stop()
+		}
+		// best-effort remove temp file
+		_ = os.Remove(cnsJsonFileName)
+		// clear globals to avoid cross-test leakage
+		svc = nil
+		service = nil
+		mux = nil
+	}
+
+	return stop, perTestSvc, nil
 }
 
 func contains(networkContainers []cns.GetNetworkContainerResponse, str string) bool {
@@ -1918,8 +1939,18 @@ func TestGetVMUniqueIDFailed(t *testing.T) {
 }
 
 func TestIBDevices(t *testing.T) {
+	// If a particular test needs to run with an isolated ServiceConfig (different ChannelMode
+	// for example), startService now returns a stop func and the instantiated *HTTPRestService.
+	// Example per-test usage (uncomment to use):
+	stop, perTestSvc, err := startService(common.ServiceConfig{ChannelMode: cns.CRD}, configuration.CNSConfig{})
+	if err != nil {
+		t.Fatalf("failed to start per-test service: %v", err)
+	}
+	defer stop()
+	// You can then use 'perTestSvc' for test-local assertions and the returned stop func will
+	// cleanly tear down the service when the test finishes.
+
 	var (
-		err          error
 		req          *http.Request
 		ibDevicesReq cns.AssignIBDevicesToPodRequest
 		body         bytes.Buffer