Cleanup, unit testing example, and automated unit testing

jpahm · jpahm · commit 4d1a1b7d8af5 · 2025-02-19T18:35:19.000-06:00
diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
@@ -33,6 +33,9 @@ jobs:
     - name: Build
       run: make build
 
+    - name: Test
+      run: make test
+
   build_windows:
     runs-on: windows-latest
     steps:
@@ -45,5 +48,11 @@ jobs:
       with:
         go-version: '1.23'
 
-    - name: Check & Build
-      run: .\build.bat
+    - name: Check
+      run: .\build.bat check
+
+    - name: Build
+      run: .\build.bat build
+
+    - name: Test
+      run: .\build.bat test
diff --git a/Makefile b/Makefile
@@ -15,5 +15,8 @@ check:
 build: ./main.go
 	go build -o $(EXEC_NAME) ./main.go
 
+test:
+	go test ./...
+
 clean: $(EXEC_NAME)
 	rm $(EXEC_NAME)
diff --git a/build.bat b/build.bat
@@ -1,17 +1,25 @@
 @echo off
-
 ::vars
 set EXEC_NAME=api-tools.exe
 
-::setup
+::simulate makefile-like branching
+set skip=1
+if "%1"=="setup" goto %1
+if "%1"=="check" goto %1
+if "%1"=="build" goto %1
+if "%1"=="test" goto %1
+set skip=0
+
+:setup
 echo Performing setup...
 go install honnef.co/go/tools/cmd/staticcheck@latest && ^
 go install golang.org/x/tools/cmd/goimports@latest
 if ERRORLEVEL 1 exit /b %ERRORLEVEL% :: fail if error occurred
 echo Setup done!
+if %skip%==1 exit
 echo[
 
-::checks
+:check
 echo Performing checks...
 go mod tidy && ^
 go vet ./... && ^
@@ -20,10 +28,21 @@ gofmt -w ./.. && ^
 goimports -w ./..
 if ERRORLEVEL 1 exit /b %ERRORLEVEL% :: fail if error occurred
 echo Checks done!
+if %skip%==1 exit
 echo[
 
-::build
+:build
 echo Building...
 go build -o %EXEC_NAME% ./main.go
 if ERRORLEVEL 1 exit /b %ERRORLEVEL% :: fail if error occurred
-echo Build complete!
+echo Build complete!
+if %skip%==1 exit
+echo[
+
+:test
+echo Testing...
+go test ./...
+if ERRORLEVEL 1 exit /b %ERRORLEVEL% :: fail if error occurred
+echo Testing complete!
+if %skip%==1 exit
+echo[
diff --git a/main.go b/main.go
@@ -11,9 +11,12 @@ import (
 	"github.com/UTDNebula/api-tools/scrapers"
 	"github.com/UTDNebula/api-tools/uploader"
 	"github.com/UTDNebula/api-tools/utils"
+	"github.com/joho/godotenv"
 )
 
 func main() {
+	// Load environment variables
+	godotenv.Load()
 
 	// Setup flags
 
diff --git a/scrapers/astra.go b/scrapers/astra.go
@@ -13,19 +13,12 @@ import (
 	"time"
 
 	"github.com/UTDNebula/api-tools/utils"
-	"github.com/joho/godotenv"
 	"github.com/valyala/fastjson"
 )
 
 var MAX_EVENTS_PER_DAY = 5000
 
 func ScrapeAstra(outDir string) {
-
-	// Load env vars
-	if err := godotenv.Load(); err != nil {
-		log.Panic("Error loading .env file")
-	}
-
 	// Start chromedp
 	chromedpCtx, cancel := utils.InitChromeDp()
 
diff --git a/scrapers/coursebook.go b/scrapers/coursebook.go
@@ -14,16 +14,9 @@ import (
 	"time"
 
 	"github.com/UTDNebula/api-tools/utils"
-	"github.com/joho/godotenv"
 )
 
 func ScrapeCoursebook(term string, startPrefix string, outDir string) {
-
-	// Load env vars
-	if err := godotenv.Load(); err != nil {
-		log.Panic("Error loading .env file")
-	}
-
 	// Start chromedp
 	chromedpCtx, cancel := utils.InitChromeDp()
 	defer cancel()
diff --git a/scrapers/mazevo.go b/scrapers/mazevo.go
@@ -14,22 +14,17 @@ import (
 	"os"
 	"time"
 
-	"github.com/joho/godotenv"
+	"github.com/UTDNebula/api-tools/utils"
 )
 
 func ScrapeMazevo(outDir string) {
-
-	// Load env vars
-	if err := godotenv.Load(); err != nil {
-		log.Panic("Error loading .env file")
-	}
-	apikey, present := os.LookupEnv("MAZEVO_API_KEY")
-	if !present {
-		log.Panic("MAZEVO_API_KEY is missing from .env!")
+	apikey, err := utils.GetEnv("MAZEVO_API_KEY")
+	if err != nil {
+		panic(err)
 	}
 
 	// Make output folder
-	err := os.MkdirAll(outDir, 0777)
+	err = os.MkdirAll(outDir, 0777)
 	if err != nil {
 		panic(err)
 	}
diff --git a/scrapers/organizations.go b/scrapers/organizations.go
@@ -10,7 +10,6 @@ import (
 	"encoding/base64"
 	"encoding/csv"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"io"
 	"log"
@@ -26,7 +25,6 @@ import (
 	"github.com/chromedp/cdproto/browser"
 	"github.com/chromedp/cdproto/network"
 	"github.com/chromedp/chromedp"
-	"github.com/joho/godotenv"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
@@ -46,16 +44,7 @@ var (
 
 func ScrapeOrganizations(outdir string) {
 	log.Println("Scraping SOC ...")
-	if err := godotenv.Load(); err != nil {
-		log.Panic("error loading .env file")
-	}
-
-	opts := append(chromedp.DefaultExecAllocatorOptions[:], chromedp.Flag("headless", false))
-	allocCtx, cancel := chromedp.NewExecAllocator(context.Background(), opts...)
-	defer cancel()
-
-	ctx, cancel := chromedp.NewContext(allocCtx)
-	// ensure cleanup occurs
+	ctx, cancel := utils.InitChromeDp()
 	defer cancel()
 
 	if err := loginToSoc(ctx); err != nil {
@@ -66,21 +55,13 @@ func ScrapeOrganizations(outdir string) {
 	}
 }
 
-func lookupEnvWithError(name string) (string, error) {
-	value, exists := os.LookupEnv(name)
-	if !exists {
-		return "", errors.New(name + " is missing from .env!")
-	}
-	return value, nil
-}
-
 func loginToSoc(ctx context.Context) error {
 	log.Println("Logging into SOC ...")
-	netID, err := lookupEnvWithError("LOGIN_NETID")
+	netID, err := utils.GetEnv("LOGIN_NETID")
 	if err != nil {
 		return err
 	}
-	password, err := lookupEnvWithError("LOGIN_PASSWORD")
+	password, err := utils.GetEnv("LOGIN_PASSWORD")
 	if err != nil {
 		return err
 	}
diff --git a/uploader/database.go b/uploader/database.go
@@ -12,6 +12,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/UTDNebula/api-tools/utils"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 )
@@ -47,10 +48,9 @@ func getCollection(client *mongo.Client, collectionName string) *mongo.Collectio
 }
 
 func getEnvMongoURI() string {
-	uri, exist := os.LookupEnv("MONGODB_URI")
-	if !exist {
-		log.Panic("Error loading 'MONGODB_URI' from the .env file")
-		os.Exit(1)
+	uri, err := utils.GetEnv("MONGODB_URI")
+	if err != nil {
+		panic(err)
 	}
 	return uri
 }
diff --git a/uploader/uploader.go b/uploader/uploader.go
@@ -20,7 +20,6 @@ import (
 	"go.mongodb.org/mongo-driver/mongo/options"
 
 	"github.com/UTDNebula/nebula-api/api/schema"
-	"github.com/joho/godotenv"
 )
 
 //  It's important to note that all of the files must be updated/uploaded TOGETHER!
@@ -33,12 +32,6 @@ import (
 var filesToUpload [3]string = [3]string{"courses.json", "professors.json", "sections.json"}
 
 func Upload(inDir string, replace bool) {
-
-	//Load env vars
-	if err := godotenv.Load(); err != nil {
-		log.Panic("Error loading .env file")
-	}
-
 	//Connect to mongo
 	client := connectDB()
 
diff --git a/utils/methods.go b/utils/methods.go
@@ -25,32 +25,41 @@ import (
 
 var Headless = true
 
+// Finds .env value and produces proper error if not found
+func GetEnv(name string) (string, error) {
+	value, exists := os.LookupEnv(name)
+	if !exists {
+		return "", errors.New(name + " is missing from .env!")
+	}
+	return value, nil
+}
+
 // Initializes Chrome DevTools Protocol
 func InitChromeDp() (chromedpCtx context.Context, cancelFnc context.CancelFunc) {
 	log.Printf("Initializing chromedp...")
 	if Headless {
 		chromedpCtx, cancelFnc = chromedp.NewContext(context.Background())
-		log.Printf("Initialized chromedp!")
 	} else {
 		allocCtx, _ := chromedp.NewExecAllocator(context.Background())
 		chromedpCtx, cancelFnc = chromedp.NewContext(allocCtx)
 	}
+	log.Printf("Initialized chromedp!")
 	return
 }
 
 // This function generates a fresh auth token and returns the new headers
 func RefreshToken(chromedpCtx context.Context) map[string][]string {
-	netID, present := os.LookupEnv("LOGIN_NETID")
-	if !present {
-		log.Panic("LOGIN_NETID is missing from .env!")
+	netID, err := GetEnv("LOGIN_NETID")
+	if err != nil {
+		panic(err)
 	}
-	password, present := os.LookupEnv("LOGIN_PASSWORD")
-	if !present {
-		log.Panic("LOGIN_PASSWORD is missing from .env!")
+	password, err := GetEnv("LOGIN_PASSWORD")
+	if err != nil {
+		panic(err)
 	}
 
 	VPrintf("Getting new token...")
-	_, err := chromedp.RunResponse(chromedpCtx,
+	_, err = chromedp.RunResponse(chromedpCtx,
 		chromedp.ActionFunc(func(ctx context.Context) error {
 			err := network.ClearBrowserCookies().Do(ctx)
 			return err
@@ -105,18 +114,18 @@ func RefreshToken(chromedpCtx context.Context) map[string][]string {
 // This function signs into Astra
 func RefreshAstraToken(chromedpCtx context.Context) map[string][]string {
 	// Get username and password
-	username, present := os.LookupEnv("LOGIN_ASTRA_USERNAME")
-	if !present {
+	username, err := GetEnv("LOGIN_ASTRA_USERNAME")
+	if err != nil {
 		log.Panic("LOGIN_ASTRA_USERNAME is missing from .env!")
 	}
-	password, present := os.LookupEnv("LOGIN_ASTRA_PASSWORD")
-	if !present {
+	password, err := GetEnv("LOGIN_ASTRA_PASSWORD")
+	if err != nil {
 		log.Panic("LOGIN_ASTRA_PASSWORD is missing from .env!")
 	}
 
 	// Sign in
 	VPrintf("Signing in...")
-	_, err := chromedp.RunResponse(chromedpCtx,
+	_, err = chromedp.RunResponse(chromedpCtx,
 		chromedp.ActionFunc(func(ctx context.Context) error {
 			err := network.ClearBrowserCookies().Do(ctx)
 			return err
diff --git a/utils/utils_test.go b/utils/utils_test.go
@@ -0,0 +1,49 @@
+package utils
+
+import (
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/joho/godotenv"
+)
+
+func TestMain(m *testing.M) {
+	// Load .env vars for testing
+	godotenv.Load("../.env")
+	// Run tests
+	os.Exit(m.Run())
+}
+
+func TestInitChromeDp(t *testing.T) {
+	// Test with head
+	Headless = false
+	ctx, cancel := InitChromeDp()
+	if ctx.Err() != nil {
+		t.Errorf("Failed to initialize chromedp with head")
+	}
+	cancel()
+	// Test headless
+	Headless = true
+	ctx, cancel = InitChromeDp()
+	if ctx.Err() != nil {
+		t.Errorf("Failed to initialize chromedp headlessly")
+	}
+	cancel()
+}
+
+func TestRefreshToken(t *testing.T) {
+	// Get a chromedp context
+	ctx, cancel := InitChromeDp()
+	defer cancel()
+	// Try refreshing token
+	headers := RefreshToken(ctx)
+	// Make sure we successfully got a PTGSESSID cookie
+	for _, cookie := range headers["Cookie"] {
+		if strings.HasPrefix(cookie, "PTGSESSID") {
+			return
+		}
+	}
+	// Fail if no PTGSESSID cookie found
+	t.Fatalf("Failed to get PTGSESSID cookie from RefreshToken!")
+}

Original file line number	Diff line number	Diff line change
`@@ -12,6 +12,7 @@ import (`
`12`	`12`	`"os"`
`13`	`13`	`"time"`
`14`	`14`
	`15`	`+ "github.com/UTDNebula/api-tools/utils"`
`15`	`16`	`"go.mongodb.org/mongo-driver/mongo"`
`16`	`17`	`"go.mongodb.org/mongo-driver/mongo/options"`
`17`	`18`	`)`
`@@ -47,10 +48,9 @@ func getCollection(client mongo.Client, collectionName string) mongo.Collectio`
`47`	`48`	`}`
`48`	`49`
`49`	`50`	`func getEnvMongoURI() string {`
`50`		`- uri, exist := os.LookupEnv("MONGODB_URI")`
`51`		`- if !exist {`
`52`		`- log.Panic("Error loading 'MONGODB_URI' from the .env file")`
`53`		`- os.Exit(1)`
	`51`	`+ uri, err := utils.GetEnv("MONGODB_URI")`
	`52`	`+ if err != nil {`
	`53`	`+ panic(err)`
`54`	`54`	`}`
`55`	`55`	`return uri`
`56`	`56`	`}`