diff --git a/src/national.py b/src/national.py
index acabbe3..b9c0d5f 100644
--- a/src/national.py
+++ b/src/national.py
@@ -41,9 +41,6 @@
 # Initialize Elexon API client
 api_client = ApiClient()
 elexon_forecast_api = GenerationForecastApi(api_client)
-elexon_forecast_generation_wind_and_solar_day_ahead_get = (
-    elexon_forecast_api.forecast_generation_wind_and_solar_day_ahead_get
-)
 
 
 @router.get(
@@ -238,7 +235,7 @@ def get_elexon_forecast(
     """
 
     try:
-        response = elexon_forecast_generation_wind_and_solar_day_ahead_get(
+        response = elexon_forecast_api.forecast_generation_wind_and_solar_day_ahead_get(
             _from=start_datetime_utc.isoformat(),
             to=end_datetime_utc.isoformat(),
             process_type=process_type,
diff --git a/src/tests/test_elexon_forecast.py b/src/tests/test_elexon_forecast.py
index 7721420..c0e4f25 100644
--- a/src/tests/test_elexon_forecast.py
+++ b/src/tests/test_elexon_forecast.py
@@ -1,14 +1,16 @@
 from typing import Optional
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pandas as pd
 import pytest
+from elexonpy.api.generation_forecast_api import GenerationForecastApi
 
 from pydantic_models import BaseModel, SolarForecastResponse
 
+API_URL = "/v0/solar/GB/national/elexon"
 
-class MockClass(BaseModel):
 
+class MockClass(BaseModel):
     start_time: str
     quantity: float
     business_type: Optional[str] = "Solar generation"
@@ -31,18 +33,17 @@ def to_dict(self):
         }
     ),
 ]
+mock_response = MagicMock()
+mock_response.data = mock_data
 
 
-class MockResponse:
-    def __init__(self):
-        self.data = mock_data
-
-
-@patch("national.elexon_forecast_generation_wind_and_solar_day_ahead_get")
+@patch.object(GenerationForecastApi, "forecast_generation_wind_and_solar_day_ahead_get")
 def test_get_elexon_forecast_mock(mock_function, api_client):
-    mock_function.return_value = MockResponse()
+    # Set mock_response
+    mock_function.return_value = mock_response
 
-    response = api_client.get("/v0/solar/GB/national/elexon")
+    # Call the API endpoint
+    response = api_client.get(API_URL)
     print("Response Headers:", response.headers)
     # Assertions
     assert response.status_code == 200
@@ -57,8 +58,7 @@ def test_get_elexon_forecast_mock(mock_function, api_client):
 
 @pytest.mark.integration
 def test_get_elexon_forecast(api_client):
-
-    response = api_client.get("/v0/solar/GB/national/elexon")
+    response = api_client.get(API_URL)
 
     # Assertions
     assert response.status_code == 200