Fix stopping of notification query on error

+ minor test fixes just btw

Author: yalegko

decoder_test.go

@@ -49,10 +49,8 @@ func TestDecoder_Unmarshal(t *testing.T) {
 	// Time pointer:
 	if system.CreationDate == nil {
 		t.Errorf("Failed to fetch Win32_Process.CreationDate")
-	} else {
-		if !(system.CreationDate.After(someOldDate) && system.CreationDate.Before(testStart)) {
-			t.Errorf("Unexoected System process creation date; %s", system.CreationDate)
-		}
+	} else if !(system.CreationDate.After(someOldDate) && system.CreationDate.Before(testStart)) {
+		t.Errorf("Unexoected System process creation date; %s", system.CreationDate)
 	}
 	// uint64:
 	if system.KernelModeTime == 0 {
@@ -242,7 +240,7 @@ func TestDecoder_Unmarshal_Slices(t *testing.T) {
 	}
 	if len(bios.BIOSVersion) < 1 {
 		t.Errorf("Empty BIOS versions list")
-	} else if len(bios.BIOSVersion[0]) < 0 {
+	} else if len(bios.BIOSVersion[0]) < 1 {
 		t.Errorf("Empty string in BIOS version liss; %v", bios.Version)
 	}
 	if len(bios.BiosCharacteristics) < 1 {
@@ -319,7 +317,9 @@ func TestDecoder_Unmarshal_References(t *testing.T) {
 	// Extract user profile of the current user.
 	var users []loggedUser
 	if err := Query(`SELECT * FROM Win32_LoggedOnUser`, &users); err != nil {
-		t.Fatalf("Failed to query Win32_LoggedOnUser; %v", err)
+		// Sometimes we can't fetch full info about all users but for the current test it's
+		// pretty ok.
+		t.Logf("Got errors querying Win32_LoggedOnUser; %v", err)
 	}
 	if len(users) < 1 {
 		t.Fatalf("No logged users found")

notification_query.go

@@ -53,7 +53,6 @@ func NewNotificationQuery(eventCh interface{}, query string) (*NotificationQuery
 	}
 	q := NotificationQuery{
 		state:   stateNotStarted,
-		doneCh:  make(chan struct{}),
 		eventCh: eventCh,
 		query:   query,
 	}
@@ -106,9 +105,16 @@ func (q *NotificationQuery) StartNotifications() (err error) {
 		q.Unlock()
 		return nil
 	}
+	q.doneCh = make(chan struct{})
 	q.state = stateStarted
 	q.Unlock()
 
+	// Mark as stopped on any return.
+	defer func() {
+		q.state = stateStopped
+		close(q.doneCh)
+	}()
+
 	//  Be aware of reflections and COM usage.
 	defer func() {
 		if r := recover(); r != nil {
@@ -150,6 +156,7 @@ func (q *NotificationQuery) StartNotifications() (err error) {
 	reflectedDoneChan := reflect.ValueOf(q.doneCh)
 	reflectedResChan := reflect.ValueOf(q.eventCh)
 	eventType := reflectedResChan.Type().Elem()
+
 	for {
 		// If it is a time to stop somebody will listen on doneCh.
 		select {

notification_query_test.go

@@ -21,8 +21,8 @@ func TestNewNotificationQuery(t *testing.T) {
 		{&T{}, true},
 		{make([]T, 0), true},
 		{make([]*T, 0), true},
-		{make(map[interface{}]T, 0), true},
-		{make(map[interface{}]*T, 0), true},
+		{make(map[interface{}]T), true},
+		{make(map[interface{}]*T), true},
 	}
 	for _, test := range cases {
 		_, err := NewNotificationQuery(test.ch, "any")
@@ -68,14 +68,7 @@ func TestNotificationQuery(t *testing.T) {
 
 	// Stop the query and confirm routine is dead.
 	query.Stop()
-	done := make(chan struct{})
-	go func() {
-		wg.Wait()
-		close(done)
-	}()
-	select {
-	case <-done:
-	case <-time.After(500 * time.Millisecond):
+	if stopped := wgWaitTimeout(&wg, 500*time.Millisecond); !stopped {
 		t.Errorf("Failed to stop query in 5x NotificationTimeout's")
 	}
 
@@ -122,14 +115,7 @@ func TestNotificationQuery_StartStop(t *testing.T) {
 	// Do not get the event!
 	// Stop the query and confirm routine is dead.
 	query.Stop()
-	done := make(chan struct{})
-	go func() {
-		wg.Wait()
-		close(done)
-	}()
-	select {
-	case <-done:
-	case <-time.After(500 * time.Millisecond):
+	if stopped := wgWaitTimeout(&wg, 500*time.Millisecond); !stopped {
 		t.Errorf("Failed to stop query in 5x NotificationTimeout's")
 	}
 }
@@ -170,14 +156,69 @@ WHERE TargetInstance ISA 'Win32_LocalTime' AND TargetInstance.Hour = 25` // Shou
 
 	// Stop the query and confirm routine is dead.
 	query.Stop()
+	if stopped := wgWaitTimeout(&wg, 500*time.Millisecond); !stopped {
+		t.Errorf("Failed to stop query in 5x NotificationTimeout's")
+	}
+}
+
+func TestNotificationQuery_StopWithError(t *testing.T) {
+	// A struct with incorrect fields that can't be unmarshaled.
+	type event struct {
+		StrangeField uint64 `wmi:"me_should_not_be_in_event"`
+	}
+
+	// Create a query that will receive an event in a short time.
+	resultCh := make(chan event)
+	queryString := `
+SELECT * FROM __InstanceModificationEvent
+WHERE TargetInstance ISA 'Win32_LocalTime'`
+
+	query, err := NewNotificationQuery(resultCh, queryString)
+	if err != nil {
+		t.Fatalf("Failed to create NotificationQuery; %s", err)
+	}
+	query.SetNotificationTimeout(20 * time.Millisecond)
+	const deadline = 1500 * time.Millisecond
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		if err := query.StartNotifications(); err == nil { // Here!
+			t.Error("Failed to report an error from StartNotifications call")
+		}
+		wg.Done()
+	}()
+
+	// Wait some time to get the error from StartNotifications about inability to
+	// parse the received object.
+	if stopped := wgWaitTimeout(&wg, deadline); !stopped {
+		t.Errorf("Failed to receive an event in %s", deadline)
+	}
+
+	// Stop the query and confirm routine is dead.
+	var stopWG sync.WaitGroup
+	stopWG.Add(1)
+	go func() {
+		query.Stop()
+		stopWG.Done()
+	}()
+	if stopped := wgWaitTimeout(&stopWG, deadline); !stopped {
+		t.Errorf("Failed to stop query in %s", deadline)
+	}
+}
+
+// Waits for wg.Wait() no more than timeout. Returns true if wg.Wait returned
+// before timeout.
+func wgWaitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
 	done := make(chan struct{})
 	go func() {
 		wg.Wait()
 		close(done)
 	}()
 	select {
 	case <-done:
-	case <-time.After(500 * time.Millisecond):
-		t.Errorf("Failed to stop query in 5x NotificationTimeout's")
+		return true
+	case <-time.After(timeout):
+		return false
 	}
 }