diff --git a/legacy/definition.go b/legacy/definition.go
index 7992ae847..8d220b552 100644
--- a/legacy/definition.go
+++ b/legacy/definition.go
@@ -677,12 +677,24 @@ func migrateRule(baseLanguage utils.Language, exitMap map[string]flows.Exit, r R
 	case "district":
 		test := stringTest{}
 		err = json.Unmarshal(r.Test.Data, &test)
-		arguments = []string{test.Test}
+		migratedState, err := MigrateTemplate(test.Test, ExtraAsFunction)
+		if err != nil {
+			return routers.Case{}, err
+		}
+		arguments = []string{migratedState}
 
 	case "ward":
 		test := wardTest{}
 		err = json.Unmarshal(r.Test.Data, &test)
-		arguments = []string{test.District, test.State}
+		migratedDistrict, err := MigrateTemplate(test.District, ExtraAsFunction)
+		if err != nil {
+			return routers.Case{}, err
+		}
+		migratedState, err := MigrateTemplate(test.State, ExtraAsFunction)
+		if err != nil {
+			return routers.Case{}, err
+		}
+		arguments = []string{migratedDistrict, migratedState}
 
 	default:
 		return routers.Case{}, fmt.Errorf("migration of '%s' tests no supported", r.Test.Type)
diff --git a/legacy/testdata/migrations/tests.json b/legacy/testdata/migrations/tests.json
index 2b4ec1108..9472d25b8 100644
--- a/legacy/testdata/migrations/tests.json
+++ b/legacy/testdata/migrations/tests.json
@@ -154,12 +154,12 @@
     {
         "legacy_test": {
             "type": "district",
-            "test": "Kigali"
+            "test": "@flow.state"
         },
         "expected_case": {
             "uuid": "d2f852ec-7b4e-457f-ae7f-f8b243c49ff5",
             "type": "has_district",
-            "arguments": ["Kigali"],
+            "arguments": ["@run.results.state"],
             "exit_uuid": "c072ecb5-0686-40ea-8ed3-898dc1349783"
         },
         "expected_localization": {}
@@ -370,13 +370,13 @@
     {
         "legacy_test": {
             "type": "ward",
-            "state": "Kigali",
-            "district": "Gasabo"
+            "state": "@flow.state",
+            "district": "@flow.district"
         },
         "expected_case": {
             "uuid": "d2f852ec-7b4e-457f-ae7f-f8b243c49ff5",
             "type": "has_ward",
-            "arguments": ["Gasabo", "Kigali"],
+            "arguments": ["@run.results.district", "@run.results.state"],
             "exit_uuid": "c072ecb5-0686-40ea-8ed3-898dc1349783"
         },
         "expected_localization": {}