diff --git a/main.go b/main.go
index 6961eed..d909240 100644
--- a/main.go
+++ b/main.go
@@ -10,9 +10,14 @@ import (
 	"strings"
 )
 
-func main() {
-	input := flag.String("input", "", "Input JSON file")
-	output := flag.String("output", "", "Output file. Uses input with extension .ndjson if not provided")
+var (
+	input  = flag.String("input", "", "Input JSON file")
+	output = flag.String("output", "", "Output file. Uses input with extension .ndjson if not provided")
+)
+
+func run() error {
+	// input := flag.String("input", "", "Input JSON file")
+	// output := flag.String("output", "", "Output file. Uses input with extension .ndjson if not provided")
 
 	flag.Usage = func() {
 		fmt.Printf("Usage of %s:\n", os.Args[0])
@@ -71,4 +76,14 @@ func main() {
 			os.Exit(1)
 		}
 	}
+
+	return nil
+}
+
+func main() {
+	flag.Parse()
+	if err := run(); err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
 }
diff --git a/main_test.go b/main_test.go
new file mode 100644
index 0000000..34900e3
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,24 @@
+package main
+
+import (
+	"flag"
+	"os"
+	"testing"
+)
+
+func TestOutputFlag(t *testing.T) {
+	// Save the original command-line arguments and restore them at the end of the test.
+	origArgs := os.Args
+	defer func() { os.Args = origArgs }()
+
+	// Set the command-line arguments for this test.
+	os.Args = []string{"cmd", "-output", "test_output"}
+
+	// Parse the flags.
+	flag.Parse()
+
+	// Check that the output flag was correctly parsed.
+	if *output != "test_output" {
+		t.Errorf("Expected output flag to be 'test_output', but got '%s'", *output)
+	}
+}