From ef4e3e7238a54a037d54cd37b21898d576d44030 Mon Sep 17 00:00:00 2001
From: Mihai Toader <mihai@sigmacomputing.com>
Date: Wed, 6 Mar 2024 10:23:26 -0800
Subject: [PATCH] Fix `-imports` handling aliased imports in source mode

---
 .../internal/tests/import_aliased/source.go   | 11 ++++
 .../tests/import_aliased/source_mock.go       | 57 +++++++++++++++++++
 mockgen/mockgen.go                            |  4 +-
 3 files changed, 70 insertions(+), 2 deletions(-)
 create mode 100644 mockgen/internal/tests/import_aliased/source.go
 create mode 100644 mockgen/internal/tests/import_aliased/source_mock.go

diff --git a/mockgen/internal/tests/import_aliased/source.go b/mockgen/internal/tests/import_aliased/source.go
new file mode 100644
index 0000000..38a963e
--- /dev/null
+++ b/mockgen/internal/tests/import_aliased/source.go
@@ -0,0 +1,11 @@
+package import_aliased
+
+import (
+	definition_alias "github.com/package/definition"
+)
+
+//go:generate mockgen -package import_aliased -destination source_mock.go -source=source.go -imports definition_alias=github.com/package/definition
+
+type S interface {
+	M(definition_alias.X)
+}
diff --git a/mockgen/internal/tests/import_aliased/source_mock.go b/mockgen/internal/tests/import_aliased/source_mock.go
new file mode 100644
index 0000000..8883e42
--- /dev/null
+++ b/mockgen/internal/tests/import_aliased/source_mock.go
@@ -0,0 +1,57 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: source.go
+//
+// Generated by this command:
+//
+//	mockgen -package import_aliased -destination source_mock.go -source=source.go -imports definition_alias=github.com/package/definition
+//
+
+// Package import_aliased is a generated GoMock package.
+package import_aliased
+
+import (
+	reflect "reflect"
+
+	definition_alias "github.com/package/definition"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockS is a mock of S interface.
+type MockS struct {
+	ctrl     *gomock.Controller
+	recorder *MockSMockRecorder
+}
+
+// MockSMockRecorder is the mock recorder for MockS.
+type MockSMockRecorder struct {
+	mock *MockS
+}
+
+// NewMockS creates a new mock instance.
+func NewMockS(ctrl *gomock.Controller) *MockS {
+	mock := &MockS{ctrl: ctrl}
+	mock.recorder = &MockSMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockS) EXPECT() *MockSMockRecorder {
+	return m.recorder
+}
+
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockS) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
+// M mocks base method.
+func (m *MockS) M(arg0 definition_alias.X) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "M", arg0)
+}
+
+// M indicates an expected call of M.
+func (mr *MockSMockRecorder) M(arg0 any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "M", reflect.TypeOf((*MockS)(nil).M), arg0)
+}
diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
index cca4467..3efa96a 100644
--- a/mockgen/mockgen.go
+++ b/mockgen/mockgen.go
@@ -376,8 +376,8 @@ func (g *generator) Generate(pkg *model.Package, outputPkgName string, outputPac
 		// try base0, base1, ...
 		pkgName := base
 
-		if _, ok := definedImports[base]; ok {
-			pkgName = definedImports[base]
+		if _, ok := definedImports[pth]; ok {
+			pkgName = definedImports[pth]
 		}
 
 		i := 0