From b9b922e776a3e0e07099e04761bbaad730f8a0c2 Mon Sep 17 00:00:00 2001
From: Marek Wolan <marek.wolan@methods.co.uk>
Date: Thu, 2 May 2024 15:51:06 +0100
Subject: [PATCH] Make ftp client only default on computers, not servers

---
 .../simulator/network/hardware/nodes/host/computer.py       | 5 +++++
 .../simulator/network/hardware/nodes/host/host_node.py      | 2 --
 tests/conftest.py                                           | 6 +++++-
 .../network/test_multi_lan_internet_example_network.py      | 2 +-
 tests/integration_tests/system/test_database_on_node.py     | 3 ++-
 5 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/primaite/simulator/network/hardware/nodes/host/computer.py b/src/primaite/simulator/network/hardware/nodes/host/computer.py
index 0b13163e..7ce64867 100644
--- a/src/primaite/simulator/network/hardware/nodes/host/computer.py
+++ b/src/primaite/simulator/network/hardware/nodes/host/computer.py
@@ -1,4 +1,7 @@
+from typing import ClassVar, Dict
+
 from primaite.simulator.network.hardware.nodes.host.host_node import HostNode
+from primaite.simulator.system.services.ftp.ftp_client import FTPClient
 
 
 class Computer(HostNode):
@@ -29,4 +32,6 @@ class Computer(HostNode):
         * Web Browser
     """
 
+    SYSTEM_SOFTWARE: ClassVar[Dict] = {**HostNode.SYSTEM_SOFTWARE, "FTPClient": FTPClient}
+
     pass
diff --git a/src/primaite/simulator/network/hardware/nodes/host/host_node.py b/src/primaite/simulator/network/hardware/nodes/host/host_node.py
index df254b1e..caea2dd7 100644
--- a/src/primaite/simulator/network/hardware/nodes/host/host_node.py
+++ b/src/primaite/simulator/network/hardware/nodes/host/host_node.py
@@ -10,7 +10,6 @@ from primaite.simulator.network.transmission.data_link_layer import Frame
 from primaite.simulator.system.applications.web_browser import WebBrowser
 from primaite.simulator.system.services.arp.arp import ARP, ARPPacket
 from primaite.simulator.system.services.dns.dns_client import DNSClient
-from primaite.simulator.system.services.ftp.ftp_client import FTPClient
 from primaite.simulator.system.services.icmp.icmp import ICMP
 from primaite.simulator.system.services.ntp.ntp_client import NTPClient
 from primaite.utils.validators import IPV4Address
@@ -301,7 +300,6 @@ class HostNode(Node):
         "HostARP": HostARP,
         "ICMP": ICMP,
         "DNSClient": DNSClient,
-        "FTPClient": FTPClient,
         "NTPClient": NTPClient,
         "WebBrowser": WebBrowser,
     }
diff --git a/tests/conftest.py b/tests/conftest.py
index 7de2bfde..a9ed42c6 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -86,7 +86,11 @@ class TestApplication(Application):
 
 @pytest.fixture(scope="function")
 def uc2_network() -> Network:
-    return arcd_uc2_network()
+    with open(PRIMAITE_PATHS.user_config_path / "example_config" / "data_manipulation.yaml") as f:
+        cfg = yaml.safe_load(f)
+    game = PrimaiteGame.from_config(cfg)
+    return game.simulation.network
+    # return arcd_uc2_network()
 
 
 @pytest.fixture(scope="function")
diff --git a/tests/integration_tests/network/test_multi_lan_internet_example_network.py b/tests/integration_tests/network/test_multi_lan_internet_example_network.py
index f56fcbf2..f6d702d8 100644
--- a/tests/integration_tests/network/test_multi_lan_internet_example_network.py
+++ b/tests/integration_tests/network/test_multi_lan_internet_example_network.py
@@ -83,7 +83,7 @@ def test_sometech_webserver_cannot_access_ftp_on_sometech_storage_server():
     some_tech_storage_srv.file_system.create_file(file_name="test.png")
 
     web_server: Server = network.get_node_by_hostname("some_tech_web_srv")
-
+    web_server.software_manager.install(FTPClient)
     web_ftp_client: FTPClient = web_server.software_manager.software["FTPClient"]
 
     assert not web_ftp_client.request_file(
diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py
index 9a396fae..8a2a9793 100644
--- a/tests/integration_tests/system/test_database_on_node.py
+++ b/tests/integration_tests/system/test_database_on_node.py
@@ -101,7 +101,7 @@ def test_database_client_native_connection_query(uc2_network):
     """Tests DB query across the network returns HTTP status 200 and date."""
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
     db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
-
+    db_client.connect()
     assert db_client.query(sql="SELECT")
     assert db_client.query(sql="INSERT")
 
@@ -222,6 +222,7 @@ def test_database_client_cannot_query_offline_database_server(uc2_network):
 
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
     db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
+    db_client.connect()
     assert len(db_client.client_connections)
 
     # Establish a new connection to the DatabaseService