Add pubsub/centralhub package.

howbazaar · howbazaar · commit 5966152d3c29 · 2016-12-05T14:35:54.000+01:00
diff --git a/dependencies.tsv b/dependencies.tsv
@@ -37,6 +37,7 @@ github.com/juju/loggo	git	3b7ece48644d35850f4ced4c2cbc2cf8413f58e0	2016-08-18T02
 github.com/juju/mempool	git	24974d6c264fe5a29716e7d56ea24c4bd904b7cc	2016-02-05T10:49:27Z
 github.com/juju/mutex	git	59c26ee163447c5c57f63ff71610d433862013de	2016-06-17T01:09:07Z
 github.com/juju/persistent-cookiejar	git	5243747bf8f2d0897f6c7a52799327dc97d585e8	2016-11-15T13:33:28Z
+github.com/juju/pubsub	git	9dcaca7eb4340dbf685aa7b3ad4cc4f8691a33d4	2016-07-28T03:00:34Z
 github.com/juju/replicaset	git	6b5becf2232ce76656ea765d8d915d41755a1513	2016-11-25T16:08:49Z
 github.com/juju/retry	git	62c62032529169c7ec02fa48f93349604c345e1f	2015-10-29T02:48:21Z
 github.com/juju/rfc	git	ebdbbdb950cd039a531d15cdc2ac2cbd94f068ee	2016-07-11T02:42:13Z
diff --git a/pubsub/centralhub/centralhub.go b/pubsub/centralhub/centralhub.go
@@ -0,0 +1,47 @@
+// Copyright 2016 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package centralhub
+
+import (
+	"github.com/juju/errors"
+	"github.com/juju/pubsub"
+	"github.com/juju/utils"
+	"gopkg.in/juju/names.v2"
+	"gopkg.in/yaml.v2"
+)
+
+// New returns a new structured hub using yaml marshalling with an origin
+// specified. The post processing ensures that the maps all have string keys
+// so they messages can be marshalled between apiservers.
+func New(origin names.MachineTag) *pubsub.StructuredHub {
+
+	return pubsub.NewStructuredHub(
+		&pubsub.StructuredHubConfig{
+			Marshaller: &yamlMarshaller{},
+			Annotations: map[string]interface{}{
+				"origin": origin.String(),
+			},
+			PostProcess: ensureStringMaps,
+		})
+}
+
+type yamlMarshaller struct{}
+
+// Marshal implements Marshaller.
+func (*yamlMarshaller) Marshal(v interface{}) ([]byte, error) {
+	return yaml.Marshal(v)
+}
+
+// Unmarshal implements Marshaller.
+func (*yamlMarshaller) Unmarshal(data []byte, v interface{}) error {
+	return yaml.Unmarshal(data, v)
+}
+
+func ensureStringMaps(in map[string]interface{}) (map[string]interface{}, error) {
+	out, err := utils.ConformYAML(in)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	return out.(map[string]interface{}), nil
+}
diff --git a/pubsub/centralhub/centralhub_test.go b/pubsub/centralhub/centralhub_test.go
@@ -0,0 +1,116 @@
+// Copyright 2016 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package centralhub_test
+
+import (
+	"time"
+
+	"github.com/juju/pubsub"
+	jc "github.com/juju/testing/checkers"
+	gc "gopkg.in/check.v1"
+	"gopkg.in/juju/names.v2"
+
+	"github.com/juju/juju/pubsub/centralhub"
+	"github.com/juju/juju/testing"
+)
+
+type CentralHubSuite struct{}
+
+var _ = gc.Suite(&CentralHubSuite{})
+
+func (*CentralHubSuite) waitForSubscribers(c *gc.C, done <-chan struct{}) {
+	select {
+	case <-done:
+	case <-time.After(testing.ShortWait):
+		c.Fatal("subscribers not finished")
+	}
+}
+
+func (s *CentralHubSuite) TestSetsOrigin(c *gc.C) {
+	hub := centralhub.New(names.NewMachineTag("42"))
+	topic := pubsub.Topic("testing")
+	var called bool
+	unsub, err := hub.Subscribe(pubsub.MatchAll, func(t pubsub.Topic, data map[string]interface{}) {
+		c.Check(t, gc.Equals, topic)
+		expected := map[string]interface{}{
+			"key":    "value",
+			"origin": "machine-42",
+		}
+		c.Check(data, jc.DeepEquals, expected)
+		called = true
+	})
+
+	c.Assert(err, jc.ErrorIsNil)
+	defer unsub.Unsubscribe()
+
+	done, err := hub.Publish(topic, map[string]interface{}{"key": "value"})
+	c.Assert(err, jc.ErrorIsNil)
+	s.waitForSubscribers(c, done)
+	c.Assert(called, jc.IsTrue)
+}
+
+type IntStruct struct {
+	Key int `json:"key"`
+}
+
+func (s *CentralHubSuite) TestYAMLMarshalling(c *gc.C) {
+	hub := centralhub.New(names.NewMachineTag("42"))
+	topic := pubsub.Topic("testing")
+	var called bool
+	unsub, err := hub.Subscribe(pubsub.MatchAll, func(t pubsub.Topic, data map[string]interface{}) {
+		c.Check(t, gc.Equals, topic)
+		expected := map[string]interface{}{
+			"key":    1234,
+			"origin": "machine-42",
+		}
+		c.Check(data, jc.DeepEquals, expected)
+		called = true
+	})
+
+	c.Assert(err, jc.ErrorIsNil)
+	defer unsub.Unsubscribe()
+
+	// With the default JSON marshalling, integers are marshalled to floats into the map.
+	done, err := hub.Publish(topic, IntStruct{1234})
+	c.Assert(err, jc.ErrorIsNil)
+	s.waitForSubscribers(c, done)
+	c.Assert(called, jc.IsTrue)
+}
+
+type NestedStruct struct {
+	Key    string    `yaml:"key"`
+	Nested IntStruct `yaml:"nested"`
+}
+
+func (s *CentralHubSuite) TestPostProcessingMaps(c *gc.C) {
+	// Due to the need to send the resulting maps over the API, nested structs
+	// need to be map[string]interface{} not map[interface{}]interface{},
+	// which is what the YAML marshaller will give us.
+	hub := centralhub.New(names.NewMachineTag("42"))
+	topic := pubsub.Topic("testing")
+	var called bool
+	unsub, err := hub.Subscribe(pubsub.MatchAll, func(t pubsub.Topic, data map[string]interface{}) {
+		c.Check(t, gc.Equals, topic)
+		expected := map[string]interface{}{
+			"key": "value",
+			"nested": map[string]interface{}{
+				"key": 1234,
+			},
+			"origin": "machine-42",
+		}
+		c.Check(data, jc.DeepEquals, expected)
+		called = true
+	})
+
+	c.Assert(err, jc.ErrorIsNil)
+	defer unsub.Unsubscribe()
+
+	// With the default JSON marshalling, integers are marshalled to floats into the map.
+	done, err := hub.Publish(topic, NestedStruct{
+		Key:    "value",
+		Nested: IntStruct{1234}})
+	c.Assert(err, jc.ErrorIsNil)
+	s.waitForSubscribers(c, done)
+	c.Assert(called, jc.IsTrue)
+}
diff --git a/pubsub/centralhub/package_test.go b/pubsub/centralhub/package_test.go
@@ -0,0 +1,14 @@
+// Copyright 2016 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package centralhub_test
+
+import (
+	"testing"
+
+	gc "gopkg.in/check.v1"
+)
+
+func TestPackage(t *testing.T) {
+	gc.TestingT(t)
+}
