diff --git a/2021/3/improved.go b/2021/3/improved.go
index f7bd0a7..eeacd39 100644
--- a/2021/3/improved.go
+++ b/2021/3/improved.go
@@ -7,11 +7,6 @@ import (
 	"strconv"
 )
 
-func main() {
-	part1()
-	part2()
-}
-
 func count(entries []string, width int) [2][]int {
 	counts := [2][]int{make([]int, width), make([]int, width)}
 
@@ -58,7 +53,7 @@ func process(input *bufio.Scanner) ([]string, int) {
 	return numbers, width
 }
 
-func part1() {
+func main() {
 	aoc := aoc.New(2021, 3)
 
 	aoc.Test(`
@@ -74,9 +69,9 @@ func part1() {
 11001
 00010
 01010
-	`, 198)
+	`, []int{198, 230})
 
-	aoc.Solution(func (input *bufio.Scanner) int {
+	aoc.Solution(1, func (input *bufio.Scanner) int {
 		entries, width := process(input)
 		counts := count(entries, width)
 
@@ -92,27 +87,8 @@ func part1() {
 
 		return gamma * epsilon
 	})
-}
 
-func part2() {
-	aoc := aoc.New(2021, 3)
-
-	aoc.Test(`
-00100
-11110
-10110
-10111
-10101
-01111
-00111
-11100
-10000
-11001
-00010
-01010
-	`, 230)
-
-	aoc.Solution(func (input *bufio.Scanner) int {
+	aoc.Solution(2, func (input *bufio.Scanner) int {
 		oxygen, width := process(input)
 		oxygenCounts := count(oxygen, width)
 
diff --git a/2021/3/main.go b/2021/3/main.go
index 32a7c00..93af264 100644
--- a/2021/3/main.go
+++ b/2021/3/main.go
@@ -6,67 +6,6 @@ import (
 	"strconv"
 )
 
-func main() {
-	part1()
-	part2()
-}
-
-func part1() {
-	aoc := aoc.New(2021, 3)
-
-	aoc.Test(`
-00100
-11110
-10110
-10111
-10101
-01111
-00111
-11100
-10000
-11001
-00010
-01010
-	`, 198)
-
-	aoc.Solution(func (input *bufio.Scanner) int {
-		length := 0
-
-		var entries []int
-
-		for input.Scan() {
-			line := input.Text()
-			if len(line) > length {
-				length = len(line)
-			}
-
-			num, err := strconv.ParseInt(line, 2, 64)
-			if err != nil {
-				panic(err)
-			}
-
-			entries = append(entries, int(num))
-		}
-
-		gamma := 0
-		epsilon := 0
-		for j := 0; j < length; j++ {
-			var count int
-			for _, num := range entries {
-				count += num >> j & 1
-			}
-
-			if float64(count)/float64(len(entries)) >= 0.5 {
-				gamma += 1 << j
-			} else {
-				epsilon += 1 << j
-			}
-		}
-
-		return gamma*epsilon
-	})
-}
-
 func filter(entries []int, length int, findCommon bool) int {
 	e := make([]int, len(entries))
 	copy(e, entries)
@@ -105,7 +44,7 @@ func filter(entries []int, length int, findCommon bool) int {
 	return e[0]
 }
 
-func part2() {
+func main() {
 	aoc := aoc.New(2021, 3)
 
 	aoc.Test(`
@@ -121,9 +60,46 @@ func part2() {
 11001
 00010
 01010
-	`, 230)
+	`, []int{198, 230})
 
-	aoc.Solution(func (input *bufio.Scanner) int {
+	aoc.Solution(1, func (input *bufio.Scanner) int {
+		length := 0
+
+		var entries []int
+
+		for input.Scan() {
+			line := input.Text()
+			if len(line) > length {
+				length = len(line)
+			}
+
+			num, err := strconv.ParseInt(line, 2, 64)
+			if err != nil {
+				panic(err)
+			}
+
+			entries = append(entries, int(num))
+		}
+
+		gamma := 0
+		epsilon := 0
+		for j := 0; j < length; j++ {
+			var count int
+			for _, num := range entries {
+				count += num >> j & 1
+			}
+
+			if float64(count)/float64(len(entries)) >= 0.5 {
+				gamma += 1 << j
+			} else {
+				epsilon += 1 << j
+			}
+		}
+
+		return gamma*epsilon
+	})
+
+	aoc.Solution(2, func (input *bufio.Scanner) int {
 		length := 0
 
 		var entries []int
diff --git a/2021/common/aoc.go b/2021/common/aoc.go
index f8e7a14..2576bfd 100644
--- a/2021/common/aoc.go
+++ b/2021/common/aoc.go
@@ -14,7 +14,7 @@ import (
 
 type testCase struct {
 	input string
-	result int
+	results []int
 }
 
 type adventOfCode struct {
@@ -24,7 +24,7 @@ type adventOfCode struct {
 }
 
 func New(year int, day int) adventOfCode {
-	err := godotenv.Load()
+	err := godotenv.Load("../../.env")
 	if err != nil {
 		fmt.Println("Failed to open .env")
 	}
@@ -32,7 +32,8 @@ func New(year int, day int) adventOfCode {
 	return adventOfCode{year, day, nil}
 }
 
-func (aoc adventOfCode) Solution(f func(*bufio.Scanner) int) {
+func (aoc adventOfCode) Solution(part int, f func(*bufio.Scanner) int) {
+	fmt.Printf("AoC %d - Day %d - Part %d\n", aoc.year, aoc.day, part)
 	if len(aoc.testCases) == 0 {
 		fmt.Println("No testCases provided!")
 	}
@@ -42,11 +43,11 @@ func (aoc adventOfCode) Solution(f func(*bufio.Scanner) int) {
 		input := bufio.NewScanner(strings.NewReader(testCase.input))
 		result := f(input)
 
-		if result != testCase.result {
-			fmt.Printf("Test %d failed, expected '%d', got '%d'\n", i+1, testCase.result, result)
+		if result != testCase.results[part-1] {
+			fmt.Printf("\tTest %d failed, expected '%d', got '%d'\n", i+1, testCase.results[part-1], result)
 			failed = true
 		} else {
-			fmt.Printf("Test %d passed!\n", i+1)
+			fmt.Printf("\tTest %d passed!\n", i+1)
 		}
 	}
 	if failed {
@@ -59,11 +60,11 @@ func (aoc adventOfCode) Solution(f func(*bufio.Scanner) int) {
 	}
 
 	result := f(input)
-	fmt.Printf("Solution: %d\n", result)
+	fmt.Printf("\tSolution:\n\t\t%d\n\n", result)
 }
 
-func (aoc *adventOfCode) Test(input string, result int) {
-	aoc.testCases = append(aoc.testCases, testCase{strings.TrimSpace(input), result})
+func (aoc *adventOfCode) Test(input string, results []int) {
+	aoc.testCases = append(aoc.testCases, testCase{strings.TrimSpace(input), results})
 }
 
 func (aoc adventOfCode) getInput() (*bufio.Scanner, error) {