Repository: openconfig/goyang
Pull request #200 (2021-08-20 23:53:56)

Depreciate Modules.FindModuleByPrefix method.
This method doesn't make sense because prefixes may be duplicated.

Instead, searches based on the prefix alone should be based on the textual location of an AST node, whether it is an extension, uses, augment, deviation, type, or leafref, so use the `FindModuleByPrefix` function instead.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/42311486/badge)](https://coveralls.io/builds/42311486)

Coverage increased (+0.02%) to 81.705% when pulling **67ee544c5955a18a35e4027f19554b27abf9aef3 on deprecate-findmodulebyprefix** into **faaf64bd380bb13295e8aa2646ff1db6d937c05a on simplify-prefix**.


Patch:
diff --git a/pkg/yang/entry.go b/pkg/yang/entry.go
index 6ff3bb9..9043d8b 100644
--- a/pkg/yang/entry.go
+++ b/pkg/yang/entry.go
@@ -1249,48 +1249,18 @@ func (e *Entry) Find(name string) *Entry {
 	// If parts[0] is "" then this path started with a /
 	// and we need to find our parent.
 	if parts[0] == "" {
+		parts = parts[1:]
+		contextNode := e.Node
 		for e.Parent != nil {
 			e = e.Parent
 		}
-		parts = parts[1:]
-
-		// Since this module might use a different prefix that isn't
-		// the prefix that the module itself uses then we need to resolve
-		// the remotely-used prefix into its locally-used prefix to find it.
-		pfxMap := map[string]string{
-			// Seed the map with the local module - we use GetPrefix just
-			// in case the module is a submodule.
-			e.Node.(*Module).GetPrefix(): e.Prefix.Name,
-		}
-
-		// Add a map between the prefix used in the import statement, and
-		// the prefix that is used in the module itself.
-		for _, i := range e.Node.(*Module).Import {
-			// Resolve the module using the current module set, since we may
-			// not have populated the Module for the entry yet.
-			m, ok := e.Node.(*Module).Modules.Modules[i.Name]
-			if !ok {
-				e.addError(fmt.Errorf("cannot find a module with name %s when looking at imports in %s", i.Name, e.Path()))
-				return nil
-			}
-
-			pfxMap[i.Prefix.Name] = m.Prefix.Name
-		}
-
 		if prefix, _ := getPrefix(parts[0]); prefix != "" {
-			pfx, ok := pfxMap[prefix]
-			if !ok {
-				// This is an undefined prefix within our context, so
-				// we can't do anything about resolving it.
-				e.addError(fmt.Errorf("invalid module prefix %s within module %s, defined prefix map: %v", prefix, e.Name, pfxMap))
-				return nil
-			}
-			m, err := e.Modules().FindModuleByPrefix(pfx)
-			if err != nil {
-				e.addError(err)
+			m := module(FindModuleByPrefix(contextNode, prefix))
+			if m == nil {
+				e.addError(fmt.Errorf("cannot find module giving prefix %q within context entry %q", prefix, e.Path()))
 				return nil
 			}
-			if e.Node.(*Module) != m {
+			if m != e.Node.(*Module) {
 				e = ToEntry(m)
 			}
 		}
diff --git a/pkg/yang/identity.go b/pkg/yang/identity.go
index 8c2b649..6fd464e 100644
--- a/pkg/yang/identity.go
+++ b/pkg/yang/identity.go
@@ -26,8 +26,10 @@ import (
 // to be identified by their module and name.
 type identityDictionary struct {
 	mu sync.Mutex
-	// TODO(wenovus): The key should not be prefix:ident-name, but
-	// module-name:ident-name. Prefixes are NOT unique.
+	// dict is a global cache of identities keyed by
+	// modulename:identityname, where modulename is the full name of the
+	// module to which the identity belongs. If the identity were defined
+	// in a submodule, then the parent module name is used instead.
 	dict map[string]resolvedIdentity
 }
 
@@ -53,7 +55,7 @@ func newResolvedIdentity(m *Module, i *Identity) (string, *resolvedIdentity) {
 		Module:   m,
 		Identity: i,
 	}
-	return i.PrefixedName(), r
+	return i.modulePrefixedName(), r
 }
 
 func appendIfNotIn(ids []*Identity, chk *Identity) []*Identity {
@@ -78,7 +80,7 @@ func addChildren(r *Identity, ids []*Identity) []*Identity {
 }
 
 // findIdentityBase returns the resolved identity that is corresponds to the
-// baseStr string in the context of the Module mod.
+// baseStr string in the context of the module/submodule mod.
 func (mod *Module) findIdentityBase(baseStr string) (*resolvedIdentity, []error) {
 	var base resolvedIdentity
 	var ok bool
@@ -92,19 +94,12 @@ func (mod *Module) findIdentityBase(baseStr string) (*resolvedIdentity, []error)
 	case "", rootPrefix:
 		// This is a local identity which is defined within the current
 		// module
-		keyName := fmt.Sprintf("%s:%s", rootPrefix, baseName)
+		keyName := fmt.Sprintf("%s:%s", module(mod).Name, baseName)
 		base, ok = identities.dict[keyName]
 		if !ok {
 			errs = append(errs, fmt.Errorf("%s: can't resolve the local base %s as %s", source, baseStr, keyName))
 		}
 	default:
-		// The identity we are looking for is prefix:basename.  If
-		// we already know prefix:basename then just use it.  If not,
-		// try again within the module identified by prefix.
-		if id, ok := identities.dict[baseStr]; ok {
-			base = id
-			break
-		}
 		// This is an identity which is defined within another module
 		extmod := FindModuleByPrefix(mod, basePrefix)
 		if extmod == nil {
@@ -112,20 +107,12 @@ func (mod *Module) findIdentityBase(baseStr string) (*resolvedIdentity, []error)
 				fmt.Errorf("%s: can't find external module with prefix %s", source, basePrefix))
 			break
 		}
-
-		// Run through the identities within the remote module and find the
-		// one that matches the base that we have been specified.
-		for _, rid := range extmod.Identities() {
-			if rid.Name == baseName {
-				key := rid.PrefixedName()
-				if id, ok := identities.dict[key]; ok {
-					base = id
-				} else {
-					errs = append(errs, fmt.Errorf("%s: can't find base %s", source, baseStr))
-				}
-				break
-			}
+		// The identity we are looking for is modulename:basename.
+		if id, ok := identities.dict[fmt.Sprintf("%s:%s", module(extmod).Name, baseName)]; ok {
+			base = id
+			break
 		}
+
 		// Error if we did not find the identity that had the name specified in
 		// the module it was expected to be in.
 		if base.isEmpty() {
diff --git a/pkg/yang/lex.go b/pkg/yang/lex.go
index e7eb4e0..49a0515 100644
--- a/pkg/yang/lex.go
+++ b/pkg/yang/lex.go
@@ -145,7 +145,7 @@ func newLexer(input, path string) *lexer {
 		file:   path,
 		input:  input,
 		line:   1, // humans start with 1
-		items:  make(chan *token, 3),
+		items:  make(chan *token, maxErrors),
 		state:  lexGround,
 		errout: os.Stderr,
 	}
@@ -187,12 +187,15 @@ func (l *lexer) emitText(c code, text string) {
 	if l.debug {
 		fmt.Fprintf(os.Stderr, "%v: %q\n", c, text)
 	}
-	l.items <- &token{
+	select {
+	case l.items <- &token{
 		code: c,
 		Text: text,
 		File: l.file,
 		Line: l.sline,
 		Col:  l.scol + 1,
+	}:
+	default:
 	}
 	l.consume()
 }
@@ -330,11 +333,14 @@ func (l *lexer) ErrorfAt(line, col int, f string, v ...interface{}) {
 // If more than maxErrors are encountered, a "too many errors" message is
 // displayed and processing stops (by clearing the input).
 func (l *lexer) adderror(err []byte) {
-	if l.errcnt >= maxErrors {
+	if l.errcnt == maxErrors {
 		l.pos = 0
 		l.start = 0
 		l.input = ""
 		l.errout.Write([]byte(tooMany))
+		l.errcnt++
+		return
+	} else if l.errcnt == maxErrors+1 {
 		return
 	}
 	l.errout.Write(err)
diff --git a/pkg/yang/modules.go b/pkg/yang/modules.go
index 286fdd8..7a8c576 100644
--- a/pkg/yang/modules.go
+++ b/pkg/yang/modules.go
@@ -29,7 +29,6 @@ type Modules struct {
 	Modules    map[string]*Module // All "module" nodes
 	SubModules map[string]*Module // All "submodule" nodes
 	includes   map[*Module]bool   // Modules we have already done include on
-	byPrefix   map[string]*Module // Cache of prefix lookup
 	byNS       map[string]*Module // Cache of namespace lookup
 	typeDict   *typeDictionary    // Cache for type definitions.
 }
@@ -40,7 +39,6 @@ func NewModules() *Modules {
 		Modules:    map[string]*Module{},
 		SubModules: map[string]*Module{},
 		includes:   map[*Module]bool{},
-		byPrefix:   map[string]*Module{},
 		byNS:       map[string]*Module{},
 		typeDict:   newTypeDictionary(),
 	}
@@ -235,34 +233,6 @@ func (ms *Modules) FindModuleByNamespace(ns string) (*Module, error) {
 	return found, nil
 }
 
-// FindModuleByPrefix either returns the Module specified by prefix or returns
-// an error.
-// TODO(wenovus): This should be deprecated since prefixes are not unique among
-// modules.
-func (ms *Modules) FindModuleByPrefix(prefix string) (*Module, error) {
-	if m, ok := ms.byPrefix[prefix]; ok {
-		return m, nil
-	}
-	var found *Module
-	for _, m := range ms.Modules {
-		if m.Prefix.Name == prefix {
-			switch {
-			case m == found:
-			case found != nil:
-				return nil, fmt.Errorf("prefix %s matches two or more modules (%s, %s)", prefix, found.Name, m.Name)
-			default:
-				found = m
-			}
-		}
-	}
-	if found == nil {
-		return nil, fmt.Errorf("%s: no such prefix", prefix)
-	}
-	// Don't cache negative results because new modules could be added.
-	ms.byPrefix[prefix] = found
-	return found, nil
-}
-
 // process satisfies all include and import statements and verifies that all
 // link ref paths reference a known node.  If an import or include references
 // a [sub]module that is not already known, Process will search for a .yang
diff --git a/pkg/yang/node.go b/pkg/yang/node.go
index 9c2f8b2..ee52efe 100644
--- a/pkg/yang/node.go
+++ b/pkg/yang/node.go
@@ -96,6 +96,9 @@ func getPrefix(s string) (string, string) {
 // relative to where n was defined.  If the prefix cannot be resolved then nil
 // is returned.
 func FindModuleByPrefix(n Node, prefix string) *Module {
+	if n == nil {
+		return nil
+	}
 	mod := RootNode(n)
 
 	if prefix == "" || prefix == mod.GetPrefix() {
@@ -104,7 +107,7 @@ func FindModuleByPrefix(n Node, prefix string) *Module {
 
 	for _, i := range mod.Import {
 		if prefix == i.Prefix.Name {
-			return i.Module
+			return mod.Modules.FindModule(i)
 		}
 	}
 	return nil
@@ -149,6 +152,17 @@ func RootNode(n Node) *Module {
 	return nil
 }
 
+// module returns the Module to which n belongs. If n resides in a submodule,
+// the belonging module will be returned.
+// If n is nil or a module could not be find, nil is returned.
+func module(n Node) *Module {
+	m := RootNode(n)
+	if m.Kind() == "submodule" {
+		m = m.Modules.Modules[m.BelongsTo.Name]
+	}
+	return m
+}
+
 // NodePath returns the full path of the node from the module name.
 func NodePath(n Node) string {
 	var path string
diff --git a/pkg/yang/yang.go b/pkg/yang/yang.go
index e37bd45..6a9c4a4 100644
--- a/pkg/yang/yang.go
+++ b/pkg/yang/yang.go
@@ -835,6 +835,11 @@ func (s *Identity) PrefixedName() string {
 	return fmt.Sprintf("%s:%s", RootNode(s).GetPrefix(), s.Name)
 }
 
+// modulePrefixedName returns the module-qualified name for the identity.
+func (s *Identity) modulePrefixedName() string {
+	return fmt.Sprintf("%s:%s", module(s).Name, s.Name)
+}
+
 // IsDefined behaves the same as the implementation for Enum - it returns
 // true if an identity with the name is defined within the Values of the
 // identity


Test patch:
diff --git a/pkg/yang/entry_test.go b/pkg/yang/entry_test.go
index 12f857a..0b7f6a4 100644
--- a/pkg/yang/entry_test.go
+++ b/pkg/yang/entry_test.go
@@ -1156,9 +1156,9 @@ module defaults {
 	} {
 		tname := strings.Join(tc.path, "/")
 
-		mod, err := ms.FindModuleByPrefix("defaults")
-		if err != nil {
-			t.Fatalf("[%d_%s] module not found: %v", i, tname, err)
+		mod, ok := ms.Modules["defaults"]
+		if !ok {
+			t.Fatalf("[%d] module not found: %q", i, tname)
 		}
 		defaults := ToEntry(mod)
 		dir, err := getdir(defaults, tc.path...)
@@ -2302,6 +2302,74 @@ func TestEntryFind(t *testing.T) {
 			"/t:e/t:operation/t:input":  "/test/e/operation/input",
 			"/t:e/t:operation/t:output": "/test/e/operation/output",
 		},
+	}, {
+		name: "submodule find",
+		inModules: map[string]string{
+			"test.yang": `
+				module test {
+					prefix "t";
+					namespace "urn:t";
+
+					include test1;
+
+					leaf a { type string; }
+					leaf b { type string; }
+
+					container c { leaf d { type string; } }
+
+                    rpc rpc1 {
+                        input { leaf input1 { type string; } }
+                    }
+
+                    container e {
+                        action operation {
+                          description "action";
+                          input { leaf input1 { type string; } }
+                          output { leaf output1 { type string; } }
+                        }
+                    }
+
+				}
+			`,
+			"test1.yang": `
+				submodule test1 {
+					belongs-to test {
+						prefix "t";
+					}
+
+					leaf d { type string; }
+				}
+			`,
+		},
+		inBaseEntryPath: "/test/d",
+		wantEntryPath: map[string]string{
+			// Absolute path with no prefixes.
+			"/b": "/test/b",
+			// Relative path with no prefixes.
+			"../b": "/test/b",
+			// Absolute path with prefixes.
+			"/t:b": "/test/b",
+			// Relative path with prefixes.
+			"../t:b": "/test/b",
+			// Find within a directory.
+			"/c/d": "/test/c/d",
+			// Find within a directory specified relatively.
+			"../c/d": "/test/c/d",
+			// Find within a relative directory with prefixes.
+			"../t:c/t:d": "/test/c/d",
+			"../t:c/d":   "/test/c/d",
+			"../c/t:d":   "/test/c/d",
+			// Find within an absolute directory with prefixes.
+			"/t:c/d":                    "/test/c/d",
+			"/c/t:d":                    "/test/c/d",
+			"../t:rpc1/input":           "/test/rpc1/input",
+			"/t:rpc1/input":             "/test/rpc1/input",
+			"/t:rpc1/t:input":           "/test/rpc1/input",
+			"/t:e/operation/input":      "/test/e/operation/input",
+			"/t:e/operation/output":     "/test/e/operation/output",
+			"/t:e/t:operation/t:input":  "/test/e/operation/input",
+			"/t:e/t:operation/t:output": "/test/e/operation/output",
+		},
 	}, {
 		name: "inter-module find",
 		inModules: map[string]string{
diff --git a/pkg/yang/identity_test.go b/pkg/yang/identity_test.go
index 93716d5..98b0448 100644
--- a/pkg/yang/identity_test.go
+++ b/pkg/yang/identity_test.go
@@ -218,11 +218,7 @@ func TestIdentityExtract(t *testing.T) {
 			foundIdentity := false
 			var thisID *Identity
 			for _, ri := range identities.dict {
-				// TODO(wenbli): Use definingModule helper from ygot after it's moved to goyang.
-				moduleName := ri.Module.Name
-				if ri.Module.Kind() == "submodule" {
-					moduleName = ri.Module.BelongsTo.Name
-				}
+				moduleName := module(ri.Module).Name
 				if ri.Identity.Name == ti.name && moduleName == ti.module {
 					foundIdentity = true
 					thisID = ri.Identity
@@ -252,6 +248,43 @@ func TestIdentityExtract(t *testing.T) {
 
 // Test cases for validating that identities can be resolved correctly.
 var treeTestCases = []identityTestCase{
+	{
+		name: "tree-test-case-0: Validate identity resolution across submodules",
+		in: []inputModule{
+			{
+				name: "base.yang",
+				content: `
+				  module base {
+				    namespace "urn:base";
+				    prefix "base";
+
+				    include side;
+
+				    identity REMOTE_BASE;
+				  }
+				`},
+			{
+				name: "remote.yang",
+				content: `
+				  submodule side {
+				    belongs-to base {
+				      prefix "r";
+				    }
+
+				    identity LOCAL_REMOTE_BASE {
+				      base r:REMOTE_BASE;
+				    }
+				  }
+				`},
+		},
+		identities: []identityOut{
+			{
+				module: "base",
+				name:   "REMOTE_BASE",
+				values: []string{"LOCAL_REMOTE_BASE"},
+			},
+		},
+	},
 	{
 		name: "tree-test-case-1: Validate identity resolution across modules",
 		in: []inputModule{
@@ -263,10 +296,15 @@ var treeTestCases = []identityTestCase{
 				    prefix "base";
 
 				    import remote { prefix "r"; }
+				    import remote2 { prefix "r2"; }
 
 				    identity LOCAL_REMOTE_BASE {
 				      base r:REMOTE_BASE;
 				    }
+
+				    identity LOCAL_REMOTE_BASE2 {
+				      base r2:REMOTE_BASE2;
+				    }
 				  }
 				`},
 			{
@@ -274,11 +312,21 @@ var treeTestCases = []identityTestCase{
 				content: `
 				  module remote {
 				    namespace "urn:remote";
-				    prefix "remote";
+				    prefix "r";
 
 				    identity REMOTE_BASE;
 				  }
 				`},
+			{
+				name: "remote2.yang",
+				content: `
+				  module remote2 {
+				    namespace "urn:remote2";
+				    prefix "remote";
+
+				    identity REMOTE_BASE2;
+				  }
+				`},
 		},
 		identities: []identityOut{
 			{
@@ -286,11 +334,21 @@ var treeTestCases = []identityTestCase{
 				name:   "REMOTE_BASE",
 				values: []string{"LOCAL_REMOTE_BASE"},
 			},
+			{
+				module: "remote2",
+				name:   "REMOTE_BASE2",
+				values: []string{"LOCAL_REMOTE_BASE2"},
+			},
 			{
 				module:    "base",
 				name:      "LOCAL_REMOTE_BASE",
 				baseNames: []string{"r:REMOTE_BASE"},
 			},
+			{
+				module:    "base",
+				name:      "LOCAL_REMOTE_BASE2",
+				baseNames: []string{"r2:REMOTE_BASE2"},
+			},
 		},
 	},
 	{
@@ -647,6 +705,7 @@ func TestIdentityTree(t *testing.T) {
 				if foundID == nil {
 					t.Errorf("Couldn't find identity %s in module %s", chkID.name,
 						chkID.module)
+					continue
 				}
 
 				if len(chkID.baseNames) > 0 {
diff --git a/pkg/yang/lex_test.go b/pkg/yang/lex_test.go
index d7ea06d..5d1c56a 100644
--- a/pkg/yang/lex_test.go
+++ b/pkg/yang/lex_test.go
@@ -277,6 +277,20 @@ without an ending.
 			`test.yang:1:1: missing closing */
 `,
 		},
+		{line(),
+			// Two errors too many.
+			`yang-version 1.1;description "\/\/\/\/\/\/\/\/\/\/";`,
+			9,
+			`test.yang:1:31: invalid escape sequence: \/
+test.yang:1:33: invalid escape sequence: \/
+test.yang:1:35: invalid escape sequence: \/
+test.yang:1:37: invalid escape sequence: \/
+test.yang:1:39: invalid escape sequence: \/
+test.yang:1:41: invalid escape sequence: \/
+test.yang:1:43: invalid escape sequence: \/
+test.yang:1:45: invalid escape sequence: \/
+` + tooMany,
+		},
 	} {
 		l := newLexer(tt.in, "test.yang")
 		errbuf := &bytes.Buffer{}
diff --git a/pkg/yang/modules_test.go b/pkg/yang/modules_test.go
index 004617b..b0a260c 100644
--- a/pkg/yang/modules_test.go
+++ b/pkg/yang/modules_test.go
@@ -94,40 +94,6 @@ func testModulesFindByCommonHandler(t *testing.T, i int, got, want *Module, want
 	}
 }
 
-func TestModulesFindByPrefix(t *testing.T) {
-	ms := testModulesForTestdataModulesText(t)
-
-	for i, tc := range []struct {
-		prefix    string
-		want      *Module
-		wantError string
-	}{
-		{
-			prefix:    "does-not-exist",
-			wantError: "does-not-exist: no such prefix",
-		},
-		{
-			prefix: "foo",
-			want:   ms.Modules["foo"],
-		},
-		{
-			prefix: "bar",
-			want:   ms.Modules["bar"],
-		},
-		{
-			prefix: "baz",
-			want:   ms.Modules["baz"],
-		},
-		{
-			prefix:    "duplicate",
-			wantError: "prefix duplicate matches two or more modules (dup-pre-",
-		},
-	} {
-		got, err := ms.FindModuleByPrefix(tc.prefix)
-		testModulesFindByCommonHandler(t, i, got, tc.want, tc.wantError, err)
-	}
-}
-
 func TestModulesFindByNamespace(t *testing.T) {
 	ms := testModulesForTestdataModulesText(t)
 
diff --git a/pkg/yang/node_test.go b/pkg/yang/node_test.go
index f90e181..5fdac62 100644
--- a/pkg/yang/node_test.go
+++ b/pkg/yang/node_test.go
@@ -79,9 +79,10 @@ func TestNode(t *testing.T) {
 		desc: "import reference statement",
 		inFn: func(ms *Modules) (Node, error) {
 
-			m, err := ms.FindModuleByPrefix("t")
-			if err != nil {
-				return nil, fmt.Errorf("can't find module in %v", ms)
+			const module = "test"
+			m, ok := ms.Modules[module]
+			if !ok {
+				return nil, fmt.Errorf("can't find module %q", module)
 			}
 
 			if len(m.Import) == 0 {
@@ -257,9 +258,10 @@ func TestNode(t *testing.T) {
 		desc: "import description statement",
 		inFn: func(ms *Modules) (Node, error) {
 
-			m, err := ms.FindModuleByPrefix("t")
-			if err != nil {
-				return nil, fmt.Errorf("can't find module in %v", ms)
+			const module = "test"
+			m, ok := ms.Modules[module]
+			if !ok {
+				return nil, fmt.Errorf("can't find module %q", module)
 			}
 
 			if len(m.Import) == 0 {
@@ -306,18 +308,19 @@ func TestNode(t *testing.T) {
 		desc: "Test matchingExtensions",
 		inFn: func(ms *Modules) (Node, error) {
 
-			m, err := ms.FindModuleByPrefix("t")
-			if err != nil {
-				return nil, fmt.Errorf("can't find module in %v", ms)
+			module := "test"
+			m, ok := ms.Modules[module]
+			if !ok {
+				return nil, fmt.Errorf("can't find module %q", module)
 			}
 
 			if len(m.Leaf) == 0 {
 				return nil, fmt.Errorf("node %v is missing imports", m)
 			}
 
-			_, err = ms.FindModuleByPrefix("f")
-			if err != nil {
-				return nil, err
+			module = "foo"
+			if _, ok := ms.Modules[module]; !ok {
+				return nil, fmt.Errorf("can't find module %q", module)
 			}
 
 			return m.Leaf[0].Type, nil
@@ -414,18 +417,19 @@ func TestNode(t *testing.T) {
 		desc: "Test matchingExtensions when module is not found",
 		inFn: func(ms *Modules) (Node, error) {
 
-			m, err := ms.FindModuleByPrefix("t")
-			if err != nil {
-				return nil, fmt.Errorf("can't find module in %v", ms)
+			module := "test"
+			m, ok := ms.Modules[module]
+			if !ok {
+				return nil, fmt.Errorf("can't find module %q", module)
 			}
 
 			if len(m.Leaf) == 0 {
 				return nil, fmt.Errorf("node %v is missing imports", m)
 			}
 
-			_, err = ms.FindModuleByPrefix("f")
-			if err != nil {
-				return nil, err
+			module = "foo"
+			if _, ok := ms.Modules[module]; !ok {
+				return nil, fmt.Errorf("can't find module %q", module)
 			}
 
 			return m.Leaf[0].Type, nil
@@ -523,3 +527,96 @@ func TestNode(t *testing.T) {
 		})
 	}
 }
+
+func TestModulesFindByPrefix(t *testing.T) {
+	// Some examples of where prefixes might be used are in the following
+	// YANG statements: extension, uses, augment, deviation, type, leafref.
+	// Not all are put into the test here, since the logic is the same for
+	// each.
+	modules := map[string]string{
+		"foo": `module foo { prefix "foo"; namespace "urn:foo"; include bar; leaf leafref { type leafref { path "../foo:leaf"; } } uses foo:lg; }`,
+		"bar": `submodule bar { belongs-to foo { prefix "bar"; } container c { uses bar:lg; } grouping lg { leaf leaf { type string; } } }`,
+		"baz": `module baz { prefix "foo"; namespace "urn:foo"; import foo { prefix f; } extension e; uses f:lg; foo:e; }`,
+	}
+
+	ms := NewModules()
+	for name, modtext := range modules {
+		if err := ms.Parse(modtext, name+".yang"); err != nil {
+			t.Fatalf("error parsing module %q: %v", name, err)
+		}
+	}
+	if errs := ms.Process(); errs != nil {
+		for _, err := range errs {
+			t.Errorf("error: %v", err)
+		}
+		t.Fatalf("fatal error(s) calling Process()")
+	}
+
+	for _, tt := range []struct {
+		desc   string
+		node   Node
+		prefix string
+		want   *Module
+	}{
+		{
+			desc:   "nil node",
+			node:   nil,
+			prefix: "does-not-exist",
+			want:   nil,
+		},
+		{
+			desc:   "module foo",
+			node:   ms.Modules["foo"],
+			prefix: "foo",
+			want:   ms.Modules["foo"],
+		},
+		{
+			desc:   "submodule bar",
+			node:   ms.SubModules["bar"],
+			prefix: "bar",
+			want:   ms.SubModules["bar"],
+		},
+		{
+			desc:   "module baz",
+			node:   ms.Modules["baz"],
+			prefix: "foo",
+			want:   ms.Modules["baz"],
+		},
+		{
+			desc:   "foo leafref",
+			node:   ms.Modules["foo"].Leaf[0].Type,
+			prefix: "foo",
+			want:   ms.Modules["foo"],
+		},
+		{
+			desc:   "foo uses",
+			node:   ms.Modules["foo"].Uses[0],
+			prefix: "foo",
+			want:   ms.Modules["foo"],
+		},
+		{
+			desc:   "bar uses",
+			node:   ms.SubModules["bar"].Container[0].Uses[0],
+			prefix: "bar",
+			want:   ms.SubModules["bar"],
+		},
+		{
+			desc:   "baz uses",
+			node:   ms.Modules["baz"].Uses[0],
+			prefix: "f",
+			want:   ms.Modules["foo"],
+		},
+		{
+			desc:   "baz extension",
+			node:   ms.Modules["baz"],
+			prefix: "foo",
+			want:   ms.Modules["baz"],
+		},
+	} {
+		t.Run(tt.desc, func(t *testing.T) {
+			if got := FindModuleByPrefix(tt.node, tt.prefix); got != tt.want {
+				t.Errorf("got: %v, want: %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/pkg/yang/types_test.go b/pkg/yang/types_test.go
index 7e5774b..311f17e 100644
--- a/pkg/yang/types_test.go
+++ b/pkg/yang/types_test.go
@@ -909,10 +909,12 @@ func TestTypeResolveUnions(t *testing.T) {
 	}}
 
 	getTestLeaf := func(ms *Modules) (*YangType, error) {
-		m, err := ms.FindModuleByPrefix("t")
-		if err != nil {
-			return nil, fmt.Errorf("can't find module in %v", ms)
+		const module = "test"
+		m, ok := ms.Modules[module]
+		if !ok {
+			return nil, fmt.Errorf("can't find module %q", module)
 		}
+
 		if len(m.Leaf) == 0 {
 			return nil, fmt.Errorf("node %v is missing imports", m)
 		}
@@ -1218,10 +1220,12 @@ func TestPattern(t *testing.T) {
 	}}
 
 	getTestLeaf := func(ms *Modules) (*YangType, error) {
-		m, err := ms.FindModuleByPrefix("t")
-		if err != nil {
-			return nil, fmt.Errorf("can't find module in %v", ms)
+		const module = "test"
+		m, ok := ms.Modules[module]
+		if !ok {
+			return nil, fmt.Errorf("can't find module %q", module)
 		}
+
 		if len(m.Leaf) == 0 {
 			return nil, fmt.Errorf("node %v is missing imports", m)
 		}