Add not found checks for location

app/controllers/locations_controller.rb

@@ -39,13 +39,18 @@ class LocationsController < ApplicationController
     end
   end
 
-  def edit; end
+  def edit
+    render_not_found and return unless @location && @scheme
+  end
 
-  def edit_name; end
+  def edit_name
+    render_not_found and return unless @location && @scheme
+  end
 
   def update
-    page = params[:location][:page]
+    render_not_found and return unless @location && @scheme
 
+    page = params[:location][:page]
     if @location.update(location_params)
       case page
       when "edit"
@@ -77,12 +82,12 @@ private
     @scheme = if %w[new create index edit_name].include?(action_name)
                 Scheme.find(params[:id])
               else
-                @location.scheme
+                @location&.scheme
               end
   end
 
   def find_location
-    @location = params[:location_id].present? ? Location.find(params[:location_id]) : Location.find(params[:id])
+    @location = params[:location_id].present? ? Location.find_by(id: params[:location_id]) : Location.find_by(id: params[:id])
   end
 
   def authenticate_scope!
@@ -90,7 +95,7 @@ private
   end
 
   def authenticate_action!
-    if %w[new edit update create index edit_name].include?(action_name) && !((current_user.organisation == @scheme.owning_organisation) || current_user.support?)
+    if %w[new edit update create index edit_name].include?(action_name) && !((current_user.organisation == @scheme&.owning_organisation) || current_user.support?)
       render_not_found and return
     end
   end

spec/requests/locations_controller_spec.rb

@@ -454,7 +454,7 @@ RSpec.describe LocationsController, type: :request do
         expect(page).to have_content("Add a location to this scheme")
       end
 
-      context "when trying to new location to a scheme that belongs to another organisation" do
+      context "when trying to edit a location that belongs to another organisation" do
         let(:another_scheme)  { FactoryBot.create(:scheme) }
         let(:another_location)  { FactoryBot.create(:location, scheme: another_scheme) }
 
@@ -463,6 +463,14 @@ RSpec.describe LocationsController, type: :request do
           expect(response).to have_http_status(:not_found)
         end
       end
+
+      context "when the requested location does not exist" do
+        let(:location) { OpenStruct.new(id: (Location.maximum(:id) || 0) + 1) }
+
+        it "returns not found" do
+          expect(response).to have_http_status(:not_found)
+        end
+      end
     end
 
     context "when signed in as a support user" do
@@ -635,6 +643,15 @@ RSpec.describe LocationsController, type: :request do
           expect(page).to have_content(I18n.t("activerecord.errors.models.location.attributes.type_of_unit.blank"))
         end
       end
+
+      context "when the requested location does not exist" do
+        let(:location) { OpenStruct.new(id: (Location.maximum(:id) || 0) + 1) }
+        let(:params) { {} }
+
+        it "returns not found" do
+          expect(response).to have_http_status(:not_found)
+        end
+      end
     end
 
     context "when signed in as a support user" do
@@ -1060,6 +1077,14 @@ RSpec.describe LocationsController, type: :request do
         expect(response).to have_http_status(:ok)
         expect(page).to have_content("Location name for #{location.postcode}")
       end
+
+      context "when the requested location does not exist" do
+        let(:location) { OpenStruct.new(id: (Location.maximum(:id) || 0) + 1) }
+
+        it "returns not found" do
+          expect(response).to have_http_status(:not_found)
+        end
+      end
     end
   end
 end