1 files changed, 4 insertions, 5 deletions
diff --git a/spec/models/environment_spec.rb b/spec/models/environment_spec.rb
index 87beba680d8..7b7b92a0b8d 100644
--- a/spec/models/environment_spec.rb
+++ b/spec/models/environment_spec.rb
@@ -18,6 +18,7 @@ RSpec.describe Environment, :use_clean_rails_memory_store_caching, feature_categ
 
   it { is_expected.to belong_to(:project).required }
   it { is_expected.to belong_to(:merge_request).optional }
+  it { is_expected.to belong_to(:cluster_agent).optional }
 
   it { is_expected.to have_many(:deployments) }
   it { is_expected.to have_many(:metrics_dashboard_annotations) }
@@ -1659,20 +1660,18 @@ RSpec.describe Environment, :use_clean_rails_memory_store_caching, feature_categ
     end
 
     context 'environment has a deployment' do
-      let!(:deployment) { create(:deployment, :success, environment: environment, cluster: cluster) }
-
       context 'with no cluster associated' do
-        let(:cluster) { nil }
+        let!(:deployment) { create(:deployment, :success, environment: environment) }
 
         it { is_expected.to be_nil }
       end
 
       context 'with a cluster associated' do
-        let(:cluster) { create(:cluster) }
+        let!(:deployment) { create(:deployment, :success, :on_cluster, environment: environment) }
 
         it 'calls the service finder' do
           expect(Clusters::KnativeServicesFinder).to receive(:new)
-            .with(cluster, environment).and_return(:finder)
+            .with(deployment.cluster, environment).and_return(:finder)
 
           is_expected.to eq :finder
         end