Refactor collectors

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

pkg/wmi/wmi.go

@@ -0,0 +1,89 @@
+package wmi
+
+import (
+	"bytes"
+	"fmt"
+	"reflect"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/yusufpapurcu/wmi"
+)
+
+func InitWbem(logger log.Logger) error {
+	// This initialization prevents a memory leak on WMF 5+. See
+	// https://github.com/prometheus-community/windows_exporter/issues/77 and
+	// linked issues for details.
+	_ = level.Debug(logger).Log("msg", "Initializing SWbemServices")
+	s, err := wmi.InitializeSWbemServices(wmi.DefaultClient)
+	if err != nil {
+		return err
+	}
+	wmi.DefaultClient.AllowMissingFields = true
+	wmi.DefaultClient.SWbemServicesClient = s
+
+	return nil
+}
+
+func className(src interface{}) string {
+	s := reflect.Indirect(reflect.ValueOf(src))
+	t := s.Type()
+	if s.Kind() == reflect.Slice {
+		t = t.Elem()
+	}
+	return t.Name()
+}
+
+func Query(query string, dst interface{}, connectServerArgs ...interface{}) error {
+	return wmi.Query(query, dst, connectServerArgs...)
+}
+
+func QueryNamespace(query string, dst interface{}, namespace string) error {
+	return wmi.QueryNamespace(query, dst, namespace)
+}
+
+func QueryAll(src interface{}, logger log.Logger) string {
+	var b bytes.Buffer
+	b.WriteString("SELECT * FROM ")
+	b.WriteString(className(src))
+
+	_ = level.Debug(logger).Log("msg", fmt.Sprintf("Generated WMI query %s", b.String()))
+	return b.String()
+}
+
+func QueryAllForClass(_ interface{}, class string, logger log.Logger) string {
+	var b bytes.Buffer
+	b.WriteString("SELECT * FROM ")
+	b.WriteString(class)
+
+	_ = level.Debug(logger).Log("msg", fmt.Sprintf("Generated WMI query %s", b.String()))
+	return b.String()
+}
+
+func QueryAllWhere(src interface{}, where string, logger log.Logger) string {
+	var b bytes.Buffer
+	b.WriteString("SELECT * FROM ")
+	b.WriteString(className(src))
+
+	if where != "" {
+		b.WriteString(" WHERE ")
+		b.WriteString(where)
+	}
+
+	_ = level.Debug(logger).Log("msg", fmt.Sprintf("Generated WMI query %s", b.String()))
+	return b.String()
+}
+
+func QueryAllForClassWhere(_ interface{}, class string, where string, logger log.Logger) string {
+	var b bytes.Buffer
+	b.WriteString("SELECT * FROM ")
+	b.WriteString(class)
+
+	if where != "" {
+		b.WriteString(" WHERE ")
+		b.WriteString(where)
+	}
+
+	_ = level.Debug(logger).Log("msg", fmt.Sprintf("Generated WMI query %s", b.String()))
+	return b.String()
+}

pkg/wmi/wmi_test.go

@@ -0,0 +1,117 @@
+package wmi
+
+import (
+	"testing"
+
+	"github.com/go-kit/log"
+)
+
+type fakeWmiClass struct {
+	Name         string
+	SomeProperty int
+}
+
+var (
+	mapQueryAll = func(src interface{}, class string, where string) string {
+		return QueryAll(src, log.NewNopLogger())
+	}
+	mapQueryAllWhere = func(src interface{}, class string, where string) string {
+		return QueryAllWhere(src, where, log.NewNopLogger())
+	}
+	mapQueryAllForClass = func(src interface{}, class string, where string) string {
+		return QueryAllForClass(src, class, log.NewNopLogger())
+	}
+	mapQueryAllForClassWhere = func(src interface{}, class string, where string) string {
+		return QueryAllForClassWhere(src, class, where, log.NewNopLogger())
+	}
+)
+
+type queryFunc func(src interface{}, class string, where string) string
+
+func TestCreateQuery(t *testing.T) {
+	cases := []struct {
+		desc      string
+		dst       interface{}
+		class     string
+		where     string
+		queryFunc queryFunc
+		expected  string
+	}{
+		{
+			desc:      "queryAll on single instance",
+			dst:       fakeWmiClass{},
+			queryFunc: mapQueryAll,
+			expected:  "SELECT * FROM fakeWmiClass",
+		},
+		{
+			desc:      "queryAll on slice",
+			dst:       []fakeWmiClass{},
+			queryFunc: mapQueryAll,
+			expected:  "SELECT * FROM fakeWmiClass",
+		},
+		{
+			desc:      "queryAllWhere on single instance",
+			dst:       fakeWmiClass{},
+			where:     "foo = bar",
+			queryFunc: mapQueryAllWhere,
+			expected:  "SELECT * FROM fakeWmiClass WHERE foo = bar",
+		},
+		{
+			desc:      "queryAllWhere on slice",
+			dst:       []fakeWmiClass{},
+			where:     "foo = bar",
+			queryFunc: mapQueryAllWhere,
+			expected:  "SELECT * FROM fakeWmiClass WHERE foo = bar",
+		},
+		{
+			desc:      "queryAllWhere on single instance with empty where",
+			dst:       fakeWmiClass{},
+			queryFunc: mapQueryAllWhere,
+			expected:  "SELECT * FROM fakeWmiClass",
+		},
+		{
+			desc:      "queryAllForClass on single instance",
+			dst:       fakeWmiClass{},
+			class:     "someClass",
+			queryFunc: mapQueryAllForClass,
+			expected:  "SELECT * FROM someClass",
+		},
+		{
+			desc:      "queryAllForClass on slice",
+			dst:       []fakeWmiClass{},
+			class:     "someClass",
+			queryFunc: mapQueryAllForClass,
+			expected:  "SELECT * FROM someClass",
+		},
+		{
+			desc:      "queryAllForClassWhere on single instance",
+			dst:       fakeWmiClass{},
+			class:     "someClass",
+			where:     "foo = bar",
+			queryFunc: mapQueryAllForClassWhere,
+			expected:  "SELECT * FROM someClass WHERE foo = bar",
+		},
+		{
+			desc:      "queryAllForClassWhere on slice",
+			dst:       []fakeWmiClass{},
+			class:     "someClass",
+			where:     "foo = bar",
+			queryFunc: mapQueryAllForClassWhere,
+			expected:  "SELECT * FROM someClass WHERE foo = bar",
+		},
+		{
+			desc:      "queryAllForClassWhere on single instance with empty where",
+			dst:       fakeWmiClass{},
+			class:     "someClass",
+			queryFunc: mapQueryAllForClassWhere,
+			expected:  "SELECT * FROM someClass",
+		},
+	}
+	for _, c := range cases {
+		t.Run(c.desc, func(t *testing.T) {
+			if q := c.queryFunc(c.dst, c.class, c.where); q != c.expected {
+				t.Errorf("Case %q failed: Expected %q, got %q", c.desc, c.expected, q)
+			}
+		})
+	}
+}