From bd87b9fa478ee61027ea811f87bc8300d0136c1b Mon Sep 17 00:00:00 2001
From: Supun Setunga <supun.setunga@gmail.com>
Date: Tue, 27 Aug 2024 12:24:09 -0700
Subject: [PATCH] Add test for staging with entitlements

---
 internal/migrate/staging_validator_test.go | 57 ++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/internal/migrate/staging_validator_test.go b/internal/migrate/staging_validator_test.go
index c01bbb646..a1e15dc09 100644
--- a/internal/migrate/staging_validator_test.go
+++ b/internal/migrate/staging_validator_test.go
@@ -876,4 +876,61 @@ func Test_StagingValidator(t *testing.T) {
 
 		require.NoError(t, err)
 	})
+
+	t.Run("contract update with entitlements", func(t *testing.T) {
+		// setup mocks
+		srv := setupValidatorMocks(t, []mockNetworkAccount{
+			{
+				address: flow.HexToAddress("01"),
+				contracts: map[string][]byte{"Foo": []byte(`
+				pub contract Foo {
+					pub resource Bar {}
+				}`)},
+			},
+			{
+				address: flow.HexToAddress("02"),
+				contracts: map[string][]byte{"Test": []byte(`
+				import Foo from 0x01
+				pub contract Test {
+					pub resource R {
+						pub var bar: auth &Foo.Bar?
+						init() {
+							self.bar = nil
+						}
+					}
+				}`)},
+			},
+		})
+
+		validator := newStagingValidator(srv)
+		err := validator.Validate([]stagedContractUpdate{
+			{
+				DeployLocation: simpleAddressLocation("0x01.Foo"),
+				SourceLocation: common.StringLocation("./Foo.cdc"),
+				Code: []byte(`
+				access(all) contract Foo {
+					access(all) resource Bar {
+						access(E) fun foo(){}
+					}
+					access(all) entitlement E
+				}`),
+			},
+			{
+				DeployLocation: simpleAddressLocation("0x02.Test"),
+				SourceLocation: common.StringLocation("./Test.cdc"),
+				Code: []byte(`
+				import Foo from 0x01
+				access(all) contract Test {
+					access(all) resource R {
+						access(all) var bar: auth(Foo.E) &Foo.Bar?
+						init() {
+							self.bar = nil
+						}
+					}
+				}`),
+			},
+		})
+
+		require.NoError(t, err)
+	})
 }