Classes for Connections and Devices: Composition, Context Managers, and a Retry Decorator

When a class earns its place

A dataclass holds data. A class earns its place when something needs to hold state and behavior together — a connection is a good example: it has a session that needs to be created, used across multiple calls, and reliably closed.

network_automation/connections.py — the class

import requests
from network_automation.exceptions import AuthenticationError, DeviceUnreachableError

class Connection:
  def __init__(self, base_url, token):
      self.base_url = base_url
      self.token = token
      self._session = None

  def __enter__(self):
      self._session = requests.Session()
      self._session.headers["Authorization"] = f"Bearer {self.token}"
      return self

  def __exit__(self, exc_type, exc_value, traceback):
      if self._session is not None:
          self._session.close()
      return False  # never swallow an exception

  def get_state(self, hostname, timeout=2):
      try:
          response = self._session.get(f"{self.base_url}/devices/{hostname}/vlans", timeout=timeout)
      except requests.exceptions.ConnectionError:
          raise DeviceUnreachableError(f"{hostname} did not respond")
      except requests.exceptions.Timeout:
          raise DeviceUnreachableError(f"{hostname} timed out")

      if response.status_code == 401:
          raise AuthenticationError(f"{hostname} rejected the credentials")
      if response.status_code == 503:
          raise DeviceUnreachableError(f"{hostname} is temporarily unavailable")
      response.raise_for_status()
      return response.json()["vlans"]

__enter__/__exit__ make this usable with Python's with statement. get_state() translates requests' own exceptions and status codes into this module's exception hierarchy from Lesson 3 -- one place where that translation happens, not repeated per call site.

Quick checkConcept check

Why does Connection deserve to be a class instead of just a function that opens a session each time?

Classes are always faster than functionsA connection needs to hold state (the session) across multiple calls, and guarantee it gets cleaned up — exactly the state-plus-behavior combination that earns a class its placeFunctions can't make HTTP requestsIt doesn't — a function would work identically

Cleanup that actually happens

using_connection.py

with Connection("http://127.0.0.1:8767", token="demo-token") as conn:
  vlans = conn.get_state("SW1")
  print("SW1 vlans:", vlans)

with Connection(...) as conn: calls __enter__ on the way in and __exit__ on the way out -- guaranteed, even if get_state() raises. Real, verified output.

output

SW1 vlans: [{'id': 10, 'name': 'USERS'}, {'id': 20, 'name': 'VOICE'}]

Real, verified output against this module's mock backend.

Quick checkPredict output

Inside a with Connection(...) as conn: block, conn.get_state() raises DeviceUnreachableError. Does __exit__ (session cleanup) still run?

No — an exception skips __exit__ entirelyYes — Python guarantees __exit__ runs when leaving a with block, whether that's normal completion or an exceptionOnly if the exception is caught firstOnly for AutomationError, not other exceptions

The retry loop, as a decorator

Reliable Automation wrote a bounded, selective retry loop by hand. Here it is, expressed once, reusable anywhere:

network_automation/connections.py — the decorator

import functools
import time

def retry(max_attempts=3, delay=2):
  def decorator(func):
      @functools.wraps(func)
      def wrapper(*args, **kwargs):
          for attempt in range(1, max_attempts + 1):
              try:
                  return func(*args, **kwargs)
              except DeviceUnreachableError:
                  if attempt == max_attempts:
                      raise
                  time.sleep(delay)
      return wrapper
  return decorator


@retry(max_attempts=3, delay=0.5)
def get_state_with_retry(connection, hostname):
  return connection.get_state(hostname)

The exact same policy Reliable Automation taught -- bounded attempts, retry only on DeviceUnreachableError, never on AuthenticationError -- now written once and applied with @retry(...) instead of copy-pasted into every function that needs it.

Same policy, reused -- not a new one

This is not new retry logic. It's Reliable Automation's exact policy (bounded, selective, exponential-backoff-free) expressed as a decorator instead of an inline loop, so any function can opt in with one line.

verify_retry.py

with Connection("http://127.0.0.1:8767", token="demo-token") as conn:
  vlans = get_state_with_retry(conn, "SW3")
  print("SW3 vlans (after retry):", vlans)

SW3's mock backend entry deliberately fails once (a real 503) before succeeding -- a genuine transient failure, not simulated in the lesson code.

output

SW3 vlans (after retry): [{'id': 10, 'name': 'USERS'}]

Real, verified output — the decorator genuinely retried once and succeeded on the second attempt.

Quick checkIdentify the correct code

Which line correctly applies the retry policy to a new function, get_vlan_count()?

def get_vlan_count(): retry(3)@retry(max_attempts=3, delay=0.5) def get_vlan_count(connection, hostname): ...get_vlan_count = retry(get_vlan_count)def get_vlan_count(): import retry

Quick checkScenario

A real, unreachable device (SW4) is queried through get_state_with_retry(). Does the retry decorator keep retrying forever?

Yes — retries are unboundedNo — max_attempts bounds it; after the final attempt still fails, the DeviceUnreachableError is re-raised, not swallowedNo, it retries exactly once regardless of max_attemptsIt depends on the device's IP address

Key takeaways

A class earns its place when state and behavior genuinely belong together — Connection's session is a real example, not an academic one.
__enter__/__exit__ make with Connection(...) as conn: guarantee cleanup, exception or not — verified directly.
@retry(...) expresses Reliable Automation's exact bounded, selective retry policy once, reusable on any function with one line — verified against a real transient failure (SW3) and a real unbounded-failure case (SW4).
None of this is new retry logic — it's the same policy, refactored for reuse.

What you'll learn

When a class earns its place

Cleanup that actually happens

The retry loop, as a decorator