From 9111061915102e12f0e3481c3bf88aa23a3d7324 Mon Sep 17 00:00:00 2001
From: Alex Flint <alex.flint@gmail.com>
Date: Wed, 4 Nov 2015 09:47:58 -0800
Subject: [PATCH] add tests for usage info

---
 usage.go      |  2 +-
 usage_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 usage_test.go

diff --git a/usage.go b/usage.go
index dc3ef3c..4852866 100644
--- a/usage.go
+++ b/usage.go
@@ -84,7 +84,7 @@ func (p *Parser) WriteHelp(w io.Writer) {
 			if spec.short != "" {
 				left += ", " + synopsis(spec, "-"+spec.short)
 			}
-			fmt.Print(left)
+			fmt.Fprint(w, left)
 			if spec.help != "" {
 				if len(left)+2 < colWidth {
 					fmt.Fprint(w, strings.Repeat(" ", colWidth-len(left)))
diff --git a/usage_test.go b/usage_test.go
new file mode 100644
index 0000000..6b8741d
--- /dev/null
+++ b/usage_test.go
@@ -0,0 +1,46 @@
+package arg
+
+import (
+	"bytes"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWriteUsage(t *testing.T) {
+	expectedUsage := "usage: example [--verbose] [--dataset DATASET] [--optimize OPTIMIZE] INPUT [OUTPUT [OUTPUT ...]] \n"
+
+	expectedHelp := `usage: example [--verbose] [--dataset DATASET] [--optimize OPTIMIZE] INPUT [OUTPUT [OUTPUT ...]] 
+
+positional arguments:
+  input
+  output
+
+options:
+  --verbose, -v          verbosity level
+  --dataset DATASET      dataset to use
+  --optimize OPTIMIZE, -O OPTIMIZE
+                         optimization level
+`
+	var args struct {
+		Input    string   `arg:"positional"`
+		Output   []string `arg:"positional"`
+		Verbose  bool     `arg:"-v,help:verbosity level"`
+		Dataset  string   `arg:"help:dataset to use"`
+		Optimize int      `arg:"-O,help:optimization level"`
+	}
+	p, err := NewParser(&args)
+	require.NoError(t, err)
+
+	os.Args[0] = "example"
+
+	var usage bytes.Buffer
+	p.WriteUsage(&usage)
+	assert.Equal(t, expectedUsage, usage.String())
+
+	var help bytes.Buffer
+	p.WriteHelp(&help)
+	assert.Equal(t, expectedHelp, help.String())
+}