From b714e998ccd620c0660e357e116f72676514d5e1 Mon Sep 17 00:00:00 2001
From: Jordan Ribbink <jribbink@telus.net>
Date: Tue, 13 Feb 2024 10:13:40 -0800
Subject: [PATCH] address feedback

---
 tools/analysis/analysis_test.go | 19 ++++++++-----------
 tools/analysis/programs.go      |  9 ++++++---
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/tools/analysis/analysis_test.go b/tools/analysis/analysis_test.go
index 514c5b5129..fb5d885d05 100644
--- a/tools/analysis/analysis_test.go
+++ b/tools/analysis/analysis_test.go
@@ -192,6 +192,8 @@ func TestParseError(t *testing.T) {
 			importRange ast.Range,
 		) ([]byte, error) {
 			switch location {
+			case contractLocation:
+				return []byte(contractCode), nil
 			default:
 				require.FailNowf(t,
 					"import of unknown location",
@@ -301,12 +303,12 @@ func TestHandledParserError(t *testing.T) {
 		},
 	}
 
-	require.Equal(t, 1, handlerCalls)
-
 	programs, err := analysis.Load(config, contractLocation)
 	require.NoError(t, err)
 
-	var parserError *parser.Error
+	require.Equal(t, 1, handlerCalls)
+
+	var parserError parser.Error
 	require.ErrorAs(t, programs[contractLocation].LoadError, &parserError)
 }
 
@@ -354,9 +356,8 @@ func TestHandledCheckerError(t *testing.T) {
 		},
 	}
 
-	require.Equal(t, 1, handlerCalls)
-
 	programs, err := analysis.Load(config, contractLocation)
+	require.Equal(t, 1, handlerCalls)
 	require.NoError(t, err)
 
 	var checkerError *sema.CheckerError
@@ -423,9 +424,8 @@ func TestHandledLoadErrorImportedProgram(t *testing.T) {
 		},
 	}
 
-	require.Equal(t, 2, handlerCalls)
-
 	programs, err := analysis.Load(config, contract1Location)
+	require.Equal(t, 2, handlerCalls)
 	require.NoError(t, err)
 
 	var checkerError *sema.CheckerError
@@ -434,10 +434,7 @@ func TestHandledLoadErrorImportedProgram(t *testing.T) {
 
 	// Validate that parent checker receives the imported program error despite it being handled
 	var importedProgramErr *sema.ImportedProgramError
-	loadErr := programs[contract1Location].LoadError.(analysis.ParsingCheckingError)
-	unwrapedErr := loadErr.Unwrap().(*sema.CheckerError)
-	require.Len(t, unwrapedErr.ChildErrors(), 1)
-	require.ErrorAs(t, unwrapedErr.ChildErrors()[0], &importedProgramErr)
+	require.ErrorAs(t, programs[contract1Location].LoadError, &importedProgramErr)
 }
 
 func TestStdlib(t *testing.T) {
diff --git a/tools/analysis/programs.go b/tools/analysis/programs.go
index 08801208f3..767dea6465 100644
--- a/tools/analysis/programs.go
+++ b/tools/analysis/programs.go
@@ -172,13 +172,16 @@ func (programs Programs) check(
 						return nil, err
 					}
 
+					program := programs[importedLocation]
+					checker := program.Checker
+
 					// If the imported program has a checker, use its elaboration for the import
-					if programs[importedLocation].Checker != nil {
-						elaboration = programs[importedLocation].Checker.Elaboration
+					if checker != nil {
+						elaboration = checker.Elaboration
 					}
 
 					// If the imported program had an error while loading, record it
-					loadError = programs[importedLocation].LoadError
+					loadError = program.LoadError
 				}
 
 				if loadError != nil {